From 96c00eb98ebb9f04b7b8e3b4fbf7323b44e546a5 Mon Sep 17 00:00:00 2001 From: Max Luebbering <2804731+le1nux@users.noreply.github.com> Date: Tue, 28 Jul 2026 18:41:06 +0200 Subject: [PATCH 1/4] feat(model): add Nemotron-3 Nano hybrid Mamba-Transformer MoE architecture Adds support for hybrid Mamba-Transformer architectures with sparse mixture-of-experts feed-forward layers, targeting Nemotron-3 Nano 30B-A3B (arXiv:2512.20848). Reference behaviour taken from Megatron-Bridge's HybridModelProvider recipe and Megatron-LM's megatron/core/ssm. The architecture is a sequence of single-operator pre-norm residual layers described by a pattern string (M = Mamba-2, * = attention, E = MoE, - = dense MLP) rather than a stack of attention+MLP blocks. Nemotron-3 Nano uses 52 layers of which only 6 attend; the Mamba-2 layers carry the positional information, so the model has no positional embeddings. Network components are registered as layer *specs* (builders) rather than modules. The component factory memoises config nodes, so injecting instantiated layers would make all layers of a type share one weight tensor. The model calls build() once per layer position instead, which keeps every hyperparameter configurable from YAML while giving each layer independent parameters. This mirrors Megatron-LM's ModuleSpec. The module tree deliberately matches GPT2's (transformer.{wte,h,lm_head_norm,lm_head}, with h an nn.ModuleDict), so the existing activation-checkpointing, FSDP2, pipeline-splitting and chunked-loss components apply unchanged. New components: * model/nemotron and nemotron_layer_spec/{mamba2,attention,moe,mlp} * Mamba-2 mixer with a dependency-free pure-PyTorch state space scan (ssd_backend=native) plus an optional fused Triton backend (ssd_backend=fused, new "mamba" extra) * MoE block with sigmoid top-k router, grouped-matmul experts, shared experts * optimizer/moe_load_balanced: auxiliary-loss-free load balancing (arXiv:2408.15664) as an optimizer step pre-hook, which makes the expert bias update correct under gradient accumulation * loss/{moe_aux_loss,weighted_sum} for the sequence-level balancing penalty * stages_generator/nemotron_stages_generator, weighting pipeline stages by per-layer-type cost and guaranteeing complete module coverage * mfu_calculator/nemotron, using active parameters and charging the quadratic attention term only to the attention layers * model_initialization model_type "nemotron"; the state space parameters (A_log, D, dt_bias, conv1d) are excluded from the regex-driven initializer and initialized by Mamba2Mixer.reset_parameters instead Also fixes a pre-existing bug in NamedParameterwiseNormalInitialization: it only stripped torch.compile's _orig_mod. prefix from parameter names, so any config applying activation checkpointing before weight initialization silently matched no per-layer regex and kept the model's default initialization. This affected GPT2 configs as well. 258 tests. The chunked SSD scan is validated against a step-by-step transcription of the Mamba-2 recurrence, plus chunk-size invariance, causality and state passing. Two distributed FSDP2 tests cover expert sharding and a full training step with load balancing. The GPT2 implementation is untouched. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 53 +- .../config_nemotron3_nano_30b_a3b_fsdp2.yaml | 421 ++++++++++++++ docs/components/components.md | 17 + docs/components/nemotron.md | 179 ++++++ pyproject.toml | 9 + src/modalities/config/config.py | 21 + src/modalities/config/pydantic_if_types.py | 2 + .../models/components/mamba2/__init__.py | 0 .../models/components/mamba2/mamba2_mixer.py | 351 ++++++++++++ .../models/components/mamba2/ssd.py | 335 +++++++++++ .../models/components/moe/__init__.py | 0 .../models/components/moe/experts.py | 174 ++++++ .../models/components/moe/load_balancing.py | 199 +++++++ src/modalities/models/components/moe/moe.py | 158 ++++++ .../models/components/moe/moe_losses.py | 129 +++++ .../models/components/moe/router.py | 130 +++++ src/modalities/models/components/norms.py | 58 ++ src/modalities/models/nemotron/__init__.py | 0 .../models/nemotron/layer_pattern.py | 90 +++ .../models/nemotron/nemotron_attention.py | 176 ++++++ .../models/nemotron/nemotron_layer_specs.py | 347 ++++++++++++ .../models/nemotron/nemotron_layers.py | 129 +++++ .../models/nemotron/nemotron_mlp.py | 44 ++ .../models/nemotron/nemotron_model.py | 326 +++++++++++ .../models/nemotron/nemotron_model_factory.py | 83 +++ .../nemotron/nemotron_stages_generator.py | 246 +++++++++ .../initialization_routines.py | 34 +- .../parameter_name_filters.py | 49 ++ src/modalities/registry/components.py | 43 ++ src/modalities/utils/nemotron_mfu.py | 173 ++++++ .../nemotron_fsdp2_config.yaml | 161 ++++++ .../test_nemotron_fsdp2.py | 147 +++++ tests/models/components/test_norms.py | 39 ++ tests/models/nemotron/__init__.py | 0 tests/models/nemotron/test_layer_pattern.py | 56 ++ tests/models/nemotron/test_mamba2_mixer.py | 236 ++++++++ tests/models/nemotron/test_moe.py | 440 +++++++++++++++ .../nemotron/test_moe_load_balancing.py | 340 ++++++++++++ .../nemotron/test_nemotron_attention.py | 166 ++++++ .../nemotron/test_nemotron_config_build.py | 147 +++++ .../nemotron/test_nemotron_initialization.py | 247 +++++++++ tests/models/nemotron/test_nemotron_layers.py | 141 +++++ tests/models/nemotron/test_nemotron_model.py | 522 ++++++++++++++++++ .../nemotron/test_nemotron_stages_and_mfu.py | 348 ++++++++++++ tests/models/nemotron/test_ssd.py | 267 +++++++++ .../nemotron_config_initialization.yaml | 84 +++ 46 files changed, 7313 insertions(+), 4 deletions(-) create mode 100644 config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml create mode 100644 docs/components/nemotron.md create mode 100644 src/modalities/models/components/mamba2/__init__.py create mode 100644 src/modalities/models/components/mamba2/mamba2_mixer.py create mode 100644 src/modalities/models/components/mamba2/ssd.py create mode 100644 src/modalities/models/components/moe/__init__.py create mode 100644 src/modalities/models/components/moe/experts.py create mode 100644 src/modalities/models/components/moe/load_balancing.py create mode 100644 src/modalities/models/components/moe/moe.py create mode 100644 src/modalities/models/components/moe/moe_losses.py create mode 100644 src/modalities/models/components/moe/router.py create mode 100644 src/modalities/models/components/norms.py create mode 100644 src/modalities/models/nemotron/__init__.py create mode 100644 src/modalities/models/nemotron/layer_pattern.py create mode 100644 src/modalities/models/nemotron/nemotron_attention.py create mode 100644 src/modalities/models/nemotron/nemotron_layer_specs.py create mode 100644 src/modalities/models/nemotron/nemotron_layers.py create mode 100644 src/modalities/models/nemotron/nemotron_mlp.py create mode 100644 src/modalities/models/nemotron/nemotron_model.py create mode 100644 src/modalities/models/nemotron/nemotron_model_factory.py create mode 100644 src/modalities/models/nemotron/nemotron_stages_generator.py create mode 100644 src/modalities/utils/nemotron_mfu.py create mode 100644 tests/fsdp2_parallelization/nemotron_fsdp2_config.yaml create mode 100644 tests/fsdp2_parallelization/test_nemotron_fsdp2.py create mode 100644 tests/models/components/test_norms.py create mode 100644 tests/models/nemotron/__init__.py create mode 100644 tests/models/nemotron/test_layer_pattern.py create mode 100644 tests/models/nemotron/test_mamba2_mixer.py create mode 100644 tests/models/nemotron/test_moe.py create mode 100644 tests/models/nemotron/test_moe_load_balancing.py create mode 100644 tests/models/nemotron/test_nemotron_attention.py create mode 100644 tests/models/nemotron/test_nemotron_config_build.py create mode 100644 tests/models/nemotron/test_nemotron_initialization.py create mode 100644 tests/models/nemotron/test_nemotron_layers.py create mode 100644 tests/models/nemotron/test_nemotron_model.py create mode 100644 tests/models/nemotron/test_nemotron_stages_and_mfu.py create mode 100644 tests/models/nemotron/test_ssd.py create mode 100644 tests/test_yaml_configs/nemotron_config_initialization.yaml diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 43d0c6e2d..3cab52010 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -4,6 +4,7 @@ |------------------|------------|---------------|------------------|------------------------------------------------------------------------------------------------| | [#141](#pr-141-towards-stable-modalities-version) | Bug Fix | [#129](https://github.com/Modalities/modalities/issues/129) | **Yes** | Towards stable modalities version | | [#154](pr-154-manual-swiglu-implementation) | Bug Fix | [#14](https://github.com/Modalities/modalities/issues/14) | **Yes** | Towards stable modalities version | +| [#nemotron](#pr-nemotron-3-nano-hybrid-mamba-transformer-moe-architecture) | Feature | -- | No | Nemotron-3 Nano hybrid Mamba-Transformer MoE architecture | | | | | | | @@ -217,4 +218,54 @@ This PR improves training monitoring and logging across runs besides some other * Add tutorials on Einsum Transformer (Example model integration) and profiling **Breaking Changes** -* experiments_root_path is now exposed on an API level \ No newline at end of file +* experiments_root_path is now exposed on an API level + +## PR Nemotron 3 Nano hybrid Mamba-Transformer MoE architecture + +Adds support for hybrid Mamba-Transformer architectures with sparse mixture-of-experts feed-forward +layers, targeting **Nemotron-3 Nano 30B-A3B** ([arXiv:2512.20848](https://arxiv.org/abs/2512.20848)). +See [docs/components/nemotron.md](docs/components/nemotron.md) for the full guide. + +**General changes:** +* New `model` variant `nemotron`: a decoder-only LM whose layer stack is described by a *layer + pattern* string (`M` = Mamba-2, `*` = attention, `E` = MoE, `-` = dense MLP), where each character + is one single-operator pre-norm residual layer. No positional embeddings; the Mamba-2 layers carry + the positional information. +* New `nemotron_layer_spec` component type with the variants `mamba2`, `attention`, `moe` and `mlp`. + These are *builders*, not modules: the model calls `build()` once per layer position, so repeated + layer types get independent weights while every hyperparameter stays configurable from YAML. +* New reusable components under `models/components/`: a Mamba-2 mixer with a dependency-free + pure-PyTorch state space scan (`ssd_backend: native`) plus an optional fused Triton backend + (`ssd_backend: fused`, new `mamba` extra), and a mixture-of-experts block with a sigmoid top-k + router, grouped-matmul experts and shared experts. +* New `optimizer` variant `moe_load_balanced`: an optimizer decorator implementing auxiliary-loss-free + load balancing ([arXiv:2408.15664](https://arxiv.org/abs/2408.15664)) as a step pre-hook, which + makes the expert bias update correct under gradient accumulation. +* New `loss` variants `moe_aux_loss` and `weighted_sum`, so the sequence-level MoE load-balancing + penalty can be added to the language modelling objective. +* New `stages_generator` variant `nemotron_stages_generator`, which balances pipeline stages by + per-layer-type computational cost instead of layer count, and guarantees that every module is + assigned to exactly one stage. +* New `mfu_calculator` variant `nemotron`, which uses *active* parameters and charges the quadratic + attention term only to the attention layers. +* New `model_initialization` model type `nemotron`. The state space parameters (`A_log`, `D`, + `dt_bias`, `conv1d_*`) are excluded from the regex-driven initializer and initialized by + `Mamba2Mixer.reset_parameters` instead, matching the reference distributions. +* Ready-to-run config: `config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml`. + +**Bug fixes (pre-existing, surfaced by the distributed tests for this feature):** +* `NamedParameterwiseNormalInitialization` only stripped torch.compile's `_orig_mod.` prefix from + parameter names before matching the initialization filters. Activation checkpointing and FSDP1 + insert their own segments (`_checkpoint_wrapped_module.`, `_fsdp_wrapped_module.`), so any config + that applied activation checkpointing before weight initialization silently matched *no* per-layer + regex and left the model with its default initialization. All wrapper prefixes are now normalized + away. This affected GPT2 configs as well. Note that `Llama3Initializer` (in the GPT2 package) still + has the same gap and was deliberately left untouched. + +**Breaking changes:** +* None. The GPT2 implementation is untouched. + +**Known limitations:** +* Tensor, expert and context parallelism are not supported for this architecture yet. FSDP2, + activation checkpointing, pipeline parallelism and `torch.compile` are. +* Multi-token prediction (MTP) and HuggingFace checkpoint conversion are not included. diff --git a/config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml b/config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml new file mode 100644 index 000000000..8a9186d7c --- /dev/null +++ b/config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml @@ -0,0 +1,421 @@ +# Nemotron-3 Nano 30B-A3B pretraining configuration. +# +# Architecture from the model report (arXiv:2512.20848, Table 1 / Figure 2) and cross-checked +# against the Megatron-Bridge recipe `recipes/nemotronh/h100/nemotron_3_nano.py`: +# +# 52 layers pattern MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME +# -> 23 Mamba-2, 23 MoE, 6 attention +# d_model 2688, 32 Q heads / 2 KV heads of dim 128 (note: 32*128 != 2688 by design) +# Mamba-2: 64 heads of dim 64 (d_inner 4096), state dim 128, 8 groups, conv kernel 4 +# MoE: 128 routed experts of dim 1856, top-6, 2 shared experts (fused into one MLP of 3712) +# squared ReLU, RMSNorm, no biases, no positional embeddings, untied embeddings +# 31.6B total / 3.2B active parameters +# +# NOTE ON SCALE: this config uses ssd_backend=native, which is correct but slow. For real runs +# install the optional extra (`pip install -e '.[mamba]'`) and switch to ssd_backend=fused. +# Tensor and expert parallelism are not supported for this architecture yet; see the PR notes. + +settings: + experiment_id: ${modalities_env:experiment_id} + config_file_path: ${modalities_env:config_file_path} + referencing_keys: + sample_key: input_ids + target_key: target_ids + prediction_key: logits + cuda_env: + local_rank: ${cuda_env:LOCAL_RANK} + global_rank: ${cuda_env:RANK} + world_size: ${cuda_env:WORLD_SIZE} + paths: + checkpoint_saving_path: data/checkpoints + train_dataset_path: ./data/lorem_ipsum_long.pbin + test_dataset_path: ./data/lorem_ipsum.pbin + experiments_root_path: ${modalities_env:experiments_root_path} + intervals: + training_log_interval_in_steps: 1 + checkpointing_interval_in_steps: 500 + evaluation_interval_in_steps: 500 + consistency_enforcement: + enforce_tokens_per_step_consistency: true + enforce_last_step_logged: false + enforce_last_step_evaluated: false + enforce_last_step_checkpointed: false + step_profile: + gradient_accumulation_steps: 1 + local_train_micro_batch_size: 1 + sequence_length: 8192 + dp_degree: + instance_key: dp_degree + pass_type: BY_REFERENCE + training_target: + num_target_tokens: + component_key: number_conversion + variant_key: num_tokens_from_packed_mem_map_dataset_continuous + config: + dataset_path: ${settings.paths.train_dataset_path} + sequence_length: ${settings.step_profile.sequence_length} + dp_degree: + instance_key: dp_degree + pass_type: BY_REFERENCE + local_micro_batch_size: ${settings.step_profile.local_train_micro_batch_size} + gradient_accumulation_steps: ${settings.step_profile.gradient_accumulation_steps} + num_target_steps: + component_key: number_conversion + variant_key: num_steps_from_num_tokens + config: + dp_degree: + instance_key: dp_degree + pass_type: BY_REFERENCE + local_micro_batch_size: ${settings.step_profile.local_train_micro_batch_size} + global_num_tokens: ${settings.training_target.num_target_tokens} + sequence_length: ${settings.step_profile.sequence_length} + gradient_accumulation_steps: ${settings.step_profile.gradient_accumulation_steps} + training_progress: + global_num_seen_tokens: 0 + num_seen_steps: 0 + num_seen_samples: 0 + last_step: -1 + +collate_fn: + component_key: collate_fn + variant_key: gpt_2_llm_collator + config: + sample_key: ${settings.referencing_keys.sample_key} + target_key: ${settings.referencing_keys.target_key} + +train_dataset: + component_key: dataset + variant_key: packed_mem_map_dataset_continuous + config: + raw_data_path: ${settings.paths.train_dataset_path} + sequence_length: ${settings.step_profile.sequence_length} + sample_key: ${settings.referencing_keys.sample_key} + +train_dataloader: + component_key: data_loader + variant_key: default + config: + num_workers: 2 + pin_memory: true + dataloader_tag: train + dataset: + instance_key: train_dataset + pass_type: BY_REFERENCE + batch_sampler: + component_key: batch_sampler + variant_key: default + config: + batch_size: ${settings.step_profile.local_train_micro_batch_size} + drop_last: true + sampler: + component_key: sampler + variant_key: resumable_distributed_sampler + config: + dataset: + instance_key: train_dataset + pass_type: BY_REFERENCE + rank: ${settings.cuda_env.global_rank} + num_replicas: ${settings.cuda_env.world_size} + shuffle: true + seed: 42 + drop_last: true + skip_num_global_samples: ${settings.training_progress.num_seen_samples} + collate_fn: + instance_key: collate_fn + pass_type: BY_REFERENCE + +test_dataset: + component_key: dataset + variant_key: packed_mem_map_dataset_continuous + config: + raw_data_path: ${settings.paths.test_dataset_path} + sequence_length: ${settings.step_profile.sequence_length} + sample_key: ${settings.referencing_keys.sample_key} + +test_dataloader: + component_key: data_loader + variant_key: default + config: + num_workers: 2 + pin_memory: true + dataloader_tag: test + dataset: + instance_key: test_dataset + pass_type: BY_REFERENCE + batch_sampler: + component_key: batch_sampler + variant_key: default + config: + batch_size: ${settings.step_profile.local_train_micro_batch_size} + drop_last: true + sampler: + component_key: sampler + variant_key: distributed_sampler + config: + rank: ${settings.cuda_env.global_rank} + num_replicas: ${settings.cuda_env.world_size} + shuffle: false + drop_last: true + dataset: + instance_key: test_dataset + pass_type: BY_REFERENCE + collate_fn: + instance_key: collate_fn + pass_type: BY_REFERENCE + +eval_dataloaders: + - instance_key: test_dataloader + pass_type: BY_REFERENCE + +checkpoint_saving: + component_key: checkpoint_saving + variant_key: default + config: + checkpoint_saving_strategy: + component_key: checkpoint_saving_strategy + variant_key: save_k_most_recent_checkpoints_strategy + config: + k: 3 + checkpoint_saving_execution: + component_key: checkpoint_saving_execution + variant_key: dcp + config: + checkpoint_path: ${settings.paths.checkpoint_saving_path} + global_rank: ${settings.cuda_env.global_rank} + experiment_id: ${settings.experiment_id} + +# The training objective is the language modelling loss plus the MoE load-balancing penalty. +# The penalty is computed inside each MoE layer (coefficient 1e-4, per the model report) and +# summed by the model into `moe_aux_loss`; the weighted sum below adds it with weight 1. +# The primary balancing mechanism is the auxiliary-loss-free expert bias, see the optimizer. +loss_fn: + component_key: loss + variant_key: weighted_sum + config: + weights: [1.0, 1.0] + losses: + - component_key: loss + variant_key: clm_cross_entropy_loss + config: + target_key: ${settings.referencing_keys.target_key} + prediction_key: ${settings.referencing_keys.prediction_key} + - component_key: loss + variant_key: moe_aux_loss + config: + prediction_key: moe_aux_loss + +device_mesh: + component_key: device_mesh + variant_key: default + config: + device_type: cuda + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + world_size: ${settings.cuda_env.world_size} + +dp_degree: + component_key: number_conversion + variant_key: parallel_degree + config: + device_mesh: + instance_key: device_mesh + pass_type: BY_REFERENCE + parallelism_methods: [dp_shard, dp_replicate] + +app_state: + component_key: app_state + variant_key: raw + config: + model: + instance_key: initialized_model + pass_type: BY_REFERENCE + optimizer: + instance_key: optimizer + pass_type: BY_REFERENCE + lr_scheduler: + instance_key: lr_scheduler + pass_type: BY_REFERENCE + +initialized_model: + component_key: model + variant_key: model_initialized + config: + model: + instance_key: fsdp_model + pass_type: BY_REFERENCE + model_initializer: + component_key: model_initialization + variant_key: composed + config: + model_type: nemotron + weight_init_type: scaled + mean: 0.0 + # sqrt(2 / (5 * 2688)) = 0.01725, matching the reference recipe's init_method_std. + std: auto + hidden_dim: ${model_raw.config.n_embd} + num_layers: ${model_raw.config.n_layer} + +fsdp_model: + component_key: model + variant_key: fsdp2_wrapped + config: + model: + instance_key: activation_checkpointed_model + pass_type: BY_REFERENCE + device_mesh: + instance_key: device_mesh + pass_type: BY_REFERENCE + mixed_precision_settings: + param_dtype: BF_16 + reduce_dtype: FP_32 + block_names: [Mamba2Layer, NemotronMoELayer, NemotronAttentionLayer, NemotronMLPLayer] + +activation_checkpointed_model: + component_key: model + variant_key: activation_checkpointed + config: + model: + instance_key: model_raw + pass_type: BY_REFERENCE + ac_variant: full_activation_checkpointing + layers_fqn: transformer.h + ac_fun_params: {} + +model_raw: + component_key: model + variant_key: nemotron + config: + use_meta_device: true + sample_key: ${settings.referencing_keys.sample_key} + prediction_key: ${settings.referencing_keys.prediction_key} + aux_loss_key: moe_aux_loss + sequence_length: ${settings.step_profile.sequence_length} + vocab_size: 131072 + n_embd: 2688 + n_layer: 52 + layer_pattern: "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME" + use_weight_tying: false + lm_head_norm_config: &nemotron_norm_config + norm_type: pytorch_rms_norm + config: + normalized_shape: ${model_raw.config.n_embd} + eps: 1e-5 + layer_specs: + "M": + component_key: nemotron_layer_spec + variant_key: mamba2 + config: + n_embd: ${model_raw.config.n_embd} + mamba_n_heads: 64 + mamba_head_dim: 64 + mamba_state_dim: 128 + mamba_n_groups: 8 + d_conv: 4 + chunk_size: 128 + ssd_backend: native # switch to `fused` once mamba-ssm / causal-conv1d are installed + norm_config: *nemotron_norm_config + "E": + component_key: nemotron_layer_spec + variant_key: moe + config: + n_embd: ${model_raw.config.n_embd} + num_experts: 128 + moe_ffn_hidden: 1856 + top_k: 6 + route_scale: 2.5 + score_function: sigmoid + use_expert_bias: true + router_dtype: float32 + num_shared_experts: 2 + shared_expert_ffn_hidden_per_expert: 1856 # -> one fused MLP of 3712 hidden units + aux_loss_coeff: 1.0e-4 + experts_backend: grouped_mm + norm_config: *nemotron_norm_config + "*": + component_key: nemotron_layer_spec + variant_key: attention + config: + n_embd: ${model_raw.config.n_embd} + n_head_q: 32 + n_head_kv: 2 + head_dim: 128 + attention_implementation: pytorch_flash + norm_config: *nemotron_norm_config + +lr_scheduler: + component_key: scheduler + variant_key: linear_warmup_cosine_annealing_lr + config: + optimizer: + instance_key: optimizer + pass_type: BY_REFERENCE + num_warmup_steps: 500 + num_training_steps: ${settings.training_target.num_target_steps} + max_lr: 4.5e-4 + min_lr: 4.5e-5 + +# The `moe_load_balanced` decorator adds the auxiliary-loss-free expert bias update as an optimizer +# step pre-hook. Pinning it to the optimizer step (rather than the forward pass) is what makes it +# correct under gradient accumulation: the token counts accumulate over micro-batches and are +# reduced across data-parallel ranks exactly once per step. +optimizer: + component_key: optimizer + variant_key: moe_load_balanced + config: + expert_bias_update_rate: 1.0e-3 + model: + instance_key: initialized_model + pass_type: BY_REFERENCE + device_mesh: + instance_key: device_mesh + pass_type: BY_REFERENCE + optimizer: + component_key: optimizer + variant_key: adam_w + config: + lr: 4.5e-4 + betas: [0.9, 0.95] + eps: 1e-8 + weight_decay: 0.1 + # The state space parameters (A_log, D, dt_bias, conv1d) parameterize the SSM dynamics and + # the router gate decides expert assignment; decaying either destabilizes training. + weight_decay_groups_excluded: [embedding, layernorm, ssm, router] + wrapped_model: + instance_key: initialized_model + pass_type: BY_REFERENCE + +gradient_clipper: + component_key: gradient_clipper + variant_key: fsdp2 + config: + wrapped_model: + instance_key: initialized_model + pass_type: BY_REFERENCE + norm_type: P2_NORM + max_norm: 1.0 + device_mesh: + instance_key: device_mesh + pass_type: BY_REFERENCE + +progress_subscriber: + component_key: progress_subscriber + variant_key: rich + config: + global_rank: ${settings.cuda_env.global_rank} + num_seen_steps: ${settings.training_progress.num_seen_steps} + num_target_steps: ${settings.training_target.num_target_steps} + train_dataloader_tag: ${train_dataloader.config.dataloader_tag} + eval_dataloaders: + instance_key: eval_dataloaders + pass_type: BY_REFERENCE + +evaluation_subscriber: + component_key: results_subscriber + variant_key: wandb + config: + global_rank: ${settings.cuda_env.global_rank} + project: modalities_nemotron + mode: ONLINE + experiment_id: ${settings.experiment_id} + directory: wandb_storage + config_file_path: ${settings.config_file_path} diff --git a/docs/components/components.md b/docs/components/components.md index 904023241..680f912b4 100644 --- a/docs/components/components.md +++ b/docs/components/components.md @@ -10,6 +10,20 @@ | model | fsdp_wrapped | [ModelFactory.get_fsdp_wrapped_model](../../src/modalities/models/model_factory.py) | [FSDPWrappedModelConfig](../../src/modalities/config/config.py) | [NNModel](../../src/modalities/models/model.py) | Model that has been sharded via FSDP | | model | model_initialized | [ModelFactory.get_weight_initialized_model](../../src/modalities/models/model_factory.py) | [WeightInitializedModelConfig](../../src/modalities/config/config.py) | [nn.Module](https://pytorch.org/docs/stable/generated/torch.nn.Module.html) | Model with initialized weights | | model | coca | [CoCa](../../src/modalities/models/coca/coca_model.py) | [CoCaConfig](../../src/modalities/models/coca/coca_model.py) | [NNModel](../../src/modalities/models/model.py) | [CoCa Model (Contrastive Captioners) ](https://arxiv.org/abs/2205.01917) | +| model | nemotron | [NemotronModelFactory.get_nemotron_model](../../src/modalities/models/nemotron/nemotron_model_factory.py) | [NemotronLLMConfig](../../src/modalities/models/nemotron/nemotron_model.py) | [NNModel](../../src/modalities/models/model.py) | Hybrid Mamba-Transformer with sparse MoE layers (Nemotron-3 Nano). See [nemotron.md](nemotron.md) | + +## Nemotron layer specs + +The network components of a hybrid model are configured as *layer specs*: builders that the model +invokes once per layer position, so that repeated layer types get independent weights. See +[nemotron.md](nemotron.md). + +| Component type | Component Version | Implementation | Configuration | Component Interface | Description | +|---------------------|-------------------|---------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------|-------------------------------------------------| +| nemotron_layer_spec | mamba2 | [Mamba2LayerSpec](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | [Mamba2LayerSpecConfig](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | [NemotronLayerSpecIF](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | Mamba-2 selective state space mixer layer (`M`) | +| nemotron_layer_spec | attention | [NemotronAttentionLayerSpec](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | [NemotronAttentionLayerSpecConfig](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | [NemotronLayerSpecIF](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | Grouped-query causal self-attention layer (`*`) | +| nemotron_layer_spec | moe | [NemotronMoELayerSpec](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | [NemotronMoELayerSpecConfig](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | [NemotronLayerSpecIF](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | Sparse mixture-of-experts layer (`E`) | +| nemotron_layer_spec | mlp | [NemotronMLPLayerSpec](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | [NemotronMLPLayerSpecConfig](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | [NemotronLayerSpecIF](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | Dense squared-ReLU feed-forward layer (`-`) | ## Weight Initialization @@ -24,6 +38,8 @@ The composed initializer supports seeded weight initialization for reproducibili | Component type | Component Version | Implementation | Configuration | Component Interface | Description | |----------------|------------------------|---------------------------------------------------------------|--------------------------------------------------------------------|------------------------------------------------|-----------------------------| | loss | clm_cross_entropy_loss | [CLMCrossEntropyLoss](../../src/modalities/loss_functions.py) | [CLMCrossEntropyLossConfig](../../src/modalities/config/config.py) | [Loss](../../src/modalities/loss_functions.py) | Cross-entropy loss function | +| loss | moe_aux_loss | [MoEAuxLoss](../../src/modalities/models/components/moe/moe_losses.py) | [MoEAuxLossConfig](../../src/modalities/models/components/moe/moe_losses.py) | [Loss](../../src/modalities/loss_functions.py) | Surfaces the pre-computed MoE load-balancing penalty | +| loss | weighted_sum | [WeightedSumLoss](../../src/modalities/models/components/moe/moe_losses.py) | [WeightedSumLossConfig](../../src/modalities/models/components/moe/moe_losses.py) | [Loss](../../src/modalities/loss_functions.py) | Weighted sum of several losses (e.g. CLM + MoE auxiliary loss) | ## Optimizers @@ -32,6 +48,7 @@ The composed initializer supports seeded weight initialization for reproducibili | optimizer | adam | [OptimizerFactory.get_adam](../../src/modalities/optimizers/optimizer_factory.py) | [AdamOptimizerConfig](../../src/modalities/config/config.py) | [Optimizer](../../src/modalities/models/model.py) | ADAM optimizer | | optimizer | adam_w | [OptimizerFactory.get_adam_w](../../src/modalities/optimizers/optimizer_factory.py) | [AdamWOptimizerConfig](../../src/modalities/config/config.py) | [Optimizer](../../src/modalities/models/model.py) | ADAMW Optimizer | | optimizer | checkpointed | [OptimizerFactory.get_checkpointed_optimizer](../../src/modalities/optimizers/optimizer_factory.py) | [CheckpointedOptimizerConfig](../../src/modalities/config/config.py) | [Optimizer](../../src/modalities/models/model.py) | Optimizer instantiated from checkpoint | +| optimizer | moe_load_balanced | [MoEBalancing.register_expert_bias_update_hook](../../src/modalities/models/components/moe/load_balancing.py) | [MoELoadBalancedOptimizerConfig](../../src/modalities/config/config.py) | [Optimizer](../../src/modalities/models/model.py) | Decorator adding auxiliary-loss-free MoE load balancing to any optimizer | ## LR Scheduling diff --git a/docs/components/nemotron.md b/docs/components/nemotron.md new file mode 100644 index 000000000..b63e2cffe --- /dev/null +++ b/docs/components/nemotron.md @@ -0,0 +1,179 @@ +# Nemotron hybrid Mamba-Transformer models + +Modalities supports hybrid Mamba-Transformer architectures with sparse mixture-of-experts feed-forward +layers, as introduced by NVIDIA's Nemotron-H and Nemotron-3 Nano families. The reference target is +**Nemotron-3 Nano 30B-A3B** ([arXiv:2512.20848](https://arxiv.org/abs/2512.20848)). + +A ready-to-run pretraining configuration is at +[config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml](../../config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml). + +## Architecture + +Unlike a classical transformer, whose block bundles attention and a feed-forward network, a hybrid +model is a sequence of **single-operator pre-norm residual layers**: + +``` +x = x + operator(norm(x)) +``` + +The stack is described by a **layer pattern** string with one character per layer: + +| Symbol | Layer type | +|--------|------------| +| `M` | Mamba-2 selective state space mixer | +| `*` | Grouped-query causal self-attention | +| `E` | Sparse mixture-of-experts feed-forward | +| `-` | Dense squared-ReLU feed-forward | + +Nemotron-3 Nano 30B-A3B uses 52 layers: + +``` +MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME +``` + +which is 23 Mamba-2, 23 MoE and just **6** attention layers. Because the Mamba-2 layers carry the +positional information, the model uses **no positional embeddings at all** (no RoPE, no learned +positions). The published hyperparameters are: + +| Property | Value | +|----------|-------| +| Layers / model dimension | 52 / 2688 | +| Attention | 32 Q heads, 2 KV heads, head dim **128** (note: 32·128 ≠ 2688 by design) | +| Mamba-2 | 64 heads of dim 64 (inner dim 4096), state dim 128, 8 groups, conv kernel 4, chunk 128 | +| MoE | 128 routed experts of dim 1856, top-6, 2 shared experts (one fused MLP of 3712) | +| Router | sigmoid gating, expert bias for auxiliary-loss-free balancing, route scale 2.5, fp32 | +| Activation / norm | squared ReLU (non-gated), RMSNorm, no biases | +| Embeddings | untied, vocab 131072 | +| Parameters | 31.6B total / 3.2B active per token | + +## Configuring the network components + +The components of the network are registered as **layer specs**: declarative builders rather than +instantiated modules. Modalities' component factory memoises each config node, so injecting an +instantiated layer would make every layer of that type share one weight tensor. A spec is asked to +`build()` once per layer position, giving independent parameters per layer while keeping every +hyperparameter reachable from YAML. + +```yaml +model_raw: + component_key: model + variant_key: nemotron + config: + n_embd: 2688 + n_layer: 52 + layer_pattern: "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME" + layer_specs: + "M": + component_key: nemotron_layer_spec + variant_key: mamba2 + config: {...} + "E": + component_key: nemotron_layer_spec + variant_key: moe + config: {...} + "*": + component_key: nemotron_layer_spec + variant_key: attention + config: {...} +``` + +Changing the Mamba/attention/MoE ratio is a one-line edit to `layer_pattern` (plus `n_layer`, which +is validated against it). Swapping an `E` for a `-` turns a sparse layer dense. + +## Mamba-2 kernel backends + +The Mamba-2 mixer has two interchangeable state space scan implementations, selected via +`ssd_backend`: + +| Backend | Requirements | Use | +|---------|--------------|-----| +| `native` (default) | none | Pure PyTorch, runs on CPU and GPU, `torch.compile`-friendly. Correct but noticeably slower and more activation-memory hungry. | +| `fused` | `pip install -e '.[mamba]'`, CUDA | The Triton kernels from `mamba-ssm` / `causal-conv1d`. Use this for real training runs. | + +The native backend implements the chunk-parallel block decomposition from the Mamba-2 paper and is +validated against a step-by-step transcription of the recurrence, so it is the reference the fused +path is checked against. A warning is emitted if `native` is used with `n_embd > 1024`. + +## Mixture-of-experts load balancing + +Nemotron combines two mechanisms, and Modalities exposes both: + +1. **Auxiliary-loss-free balancing** (primary). Each router keeps an additive per-expert bias that + affects expert *selection* only, never the output weights. An optimizer step pre-hook nudges the + bias of under-loaded experts up and over-loaded experts down by a fixed step size (1e-3). Enable + it by wrapping the optimizer: + + ```yaml + optimizer: + component_key: optimizer + variant_key: moe_load_balanced + config: + expert_bias_update_rate: 1.0e-3 + model: {instance_key: initialized_model, pass_type: BY_REFERENCE} + device_mesh: {instance_key: device_mesh, pass_type: BY_REFERENCE} + optimizer: + component_key: optimizer + variant_key: adam_w + config: {...} + ``` + + Attaching the update to the optimizer step (rather than the forward pass) is what makes it correct + under gradient accumulation: token counts accumulate across micro-batches and are reduced across + data-parallel ranks exactly once per step. + +2. **The classic load-balancing loss** (secondary, coefficient 1e-4). Computed per sequence inside + each MoE layer, summed by the model into the output dict under `aux_loss_key`, and added to the + language modelling loss via the `weighted_sum` loss: + + ```yaml + loss_fn: + component_key: loss + variant_key: weighted_sum + config: + weights: [1.0, 1.0] + losses: + - {component_key: loss, variant_key: clm_cross_entropy_loss, config: {...}} + - {component_key: loss, variant_key: moe_aux_loss, config: {prediction_key: moe_aux_loss}} + ``` + +## Weight initialization + +The generic, regex-driven initializer (`model_type: nemotron`) handles all linear and embedding +weights. The state space parameters (`A_log`, `D`, `dt_bias`, `conv1d_*`) follow their own +distributions, defined in `Mamba2Mixer.reset_parameters` and matching the reference implementation; +the parameter-name filters deliberately exclude them so they cannot be overwritten. Recommended +weight decay exclusions: + +```yaml +weight_decay_groups_excluded: [embedding, layernorm, ssm, router] +``` + +Decaying the SSM dynamics parameters or the router gate destabilizes training. + +## Parallelism support + +| Strategy | Status | +|----------|--------| +| FSDP2 (dp_shard / dp_replicate) | Supported. Expert weights shard along the expert dimension. | +| Activation checkpointing (all variants) | Supported via `layers_fqn: transformer.h`. | +| Pipeline parallelism | Supported via the `nemotron_stages_generator`, which balances stages by per-layer-type cost rather than layer count. | +| `torch.compile` per layer | Supported with the native backend. | +| Tensor parallelism | **Not supported yet.** Mamba's packed `in_proj` is a five-way unequal column split and the conv/`A_log`/`D`/`dt_bias` parameters need per-head sharding. | +| Expert parallelism | **Not supported.** The device mesh has no expert dimension. | +| Context parallelism | **Not supported yet.** Requires Mamba state passing across ranks. | + +## Registered components + +| Component type | Variant | Implementation | +|----------------|---------|----------------| +| `model` | `nemotron` | [NemotronModelFactory.get_nemotron_model](../../src/modalities/models/nemotron/nemotron_model_factory.py) | +| `nemotron_layer_spec` | `mamba2` | [Mamba2LayerSpec](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | +| `nemotron_layer_spec` | `attention` | [NemotronAttentionLayerSpec](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | +| `nemotron_layer_spec` | `moe` | [NemotronMoELayerSpec](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | +| `nemotron_layer_spec` | `mlp` | [NemotronMLPLayerSpec](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | +| `optimizer` | `moe_load_balanced` | [MoEBalancing.register_expert_bias_update_hook](../../src/modalities/models/components/moe/load_balancing.py) | +| `loss` | `moe_aux_loss` | [MoEAuxLoss](../../src/modalities/models/components/moe/moe_losses.py) | +| `loss` | `weighted_sum` | [WeightedSumLoss](../../src/modalities/models/components/moe/moe_losses.py) | +| `stages_generator` | `nemotron_stages_generator` | [NemotronStagesGenerator](../../src/modalities/models/nemotron/nemotron_stages_generator.py) | +| `mfu_calculator` | `nemotron` | [NemotronMFUCalculator](../../src/modalities/utils/nemotron_mfu.py) | +| `model_initialization` | `composed` (`model_type: nemotron`) | [ComposedInitializationRoutines](../../src/modalities/nn/model_initialization/composed_initialization.py) | diff --git a/pyproject.toml b/pyproject.toml index 2ddc23935..4e58fc742 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,15 @@ build-backend = "setuptools.build_meta" linting = ["pre-commit"] tests = ["pytest", "pytest-cov", "debugpy"] +# Fused Mamba-2 kernels. Optional: the hybrid Mamba-Transformer models ship a pure-PyTorch +# state space scan (ssd_backend="native") that runs everywhere. Install this extra and set +# ssd_backend="fused" for large-scale training, where the Triton kernels are substantially +# faster and use less activation memory. Requires a CUDA toolchain to build. +mamba = [ + "mamba-ssm>=2.2.0; platform_system != 'Darwin'", + "causal-conv1d>=1.4.0; platform_system != 'Darwin'", +] + cpu = ["torch>=2.10,<2.11.0", "torchvision"] cu126 = [ "torch>=2.10,<2.11.0", diff --git a/src/modalities/config/config.py b/src/modalities/config/config.py index c0bf7f2ed..0ae12f0ba 100644 --- a/src/modalities/config/config.py +++ b/src/modalities/config/config.py @@ -388,6 +388,27 @@ class SelectiveOpACParams(BaseModel): ac_fun_params: FullACParams | SelectiveLayerACParams | SelectiveOpACParams +class MoELoadBalancedOptimizerConfig(BaseModel): + """Configuration of the auxiliary-loss-free MoE load balancing optimizer decorator. + + Attributes: + optimizer (Optimizer): The optimizer to decorate. + model (nn.Module): The model whose MoE layers should be balanced. + expert_bias_update_rate (float): Step size of the per-expert bias update. Nemotron-3 Nano + uses 1e-3. + device_mesh (DeviceMesh | None): Device mesh used to resolve the reduction group for the + expert token counts. + """ + + optimizer: PydanticOptimizerIFType + model: PydanticPytorchModuleType + expert_bias_update_rate: Annotated[float, Field(strict=True, gt=0.0)] + device_mesh: Optional[PydanticDeviceMeshIFType] = None + + # avoid the pydantic warning about the protected 'model_' namespace + model_config = ConfigDict(protected_namespaces=()) + + class RawAppStateConfig(BaseModel): model: PydanticPytorchModuleOrListType optimizer: PydanticOptimizerIFType diff --git a/src/modalities/config/pydantic_if_types.py b/src/modalities/config/pydantic_if_types.py index 90b7ca951..263195340 100644 --- a/src/modalities/config/pydantic_if_types.py +++ b/src/modalities/config/pydantic_if_types.py @@ -22,6 +22,7 @@ from modalities.inference.text.inference_component import TextInferenceComponent from modalities.logging_broker.subscriber import MessageSubscriberIF from modalities.loss_functions import Loss +from modalities.models.nemotron.nemotron_layer_specs import NemotronLayerSpecIF from modalities.models.parallelism.pipeline_parallelism import Pipeline, StagesGenerator from modalities.nn.model_initialization.initialization_if import ModelInitializationIF from modalities.tokenization.tokenizer_wrapper import TokenizerWrapper @@ -98,3 +99,4 @@ def __get_pydantic_core_schema__( torch.utils.hooks.RemovableHandle, PydanticThirdPartyTypeIF(torch.utils.hooks.RemovableHandle) ] PydanticDebuggingType = Annotated[Debugging, PydanticThirdPartyTypeIF(Debugging)] +PydanticNemotronLayerSpecIFType = Annotated[NemotronLayerSpecIF, PydanticThirdPartyTypeIF(NemotronLayerSpecIF)] diff --git a/src/modalities/models/components/mamba2/__init__.py b/src/modalities/models/components/mamba2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/modalities/models/components/mamba2/mamba2_mixer.py b/src/modalities/models/components/mamba2/mamba2_mixer.py new file mode 100644 index 000000000..8fd42aaa7 --- /dev/null +++ b/src/modalities/models/components/mamba2/mamba2_mixer.py @@ -0,0 +1,351 @@ +"""Mamba-2 mixer, the sequence-mixing operator of hybrid Mamba-Transformer models. + +The layout follows the Megatron-LM reference implementation +(``megatron/core/ssm/mamba_mixer.py``) so that parameter shapes and initialization match, which +keeps checkpoint conversion between the two frameworks tractable. +""" + +import logging +import math +from enum import Enum + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed.tensor import DTensor + +from modalities.models.components.mamba2.ssd import GatedRMSNorm, causal_depthwise_conv1d, ssd_chunked_scan + +logger = logging.getLogger(__name__) + +try: + from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined +except ModuleNotFoundError: + mamba_chunk_scan_combined = None + +try: + from causal_conv1d import causal_conv1d_fn +except ModuleNotFoundError: + causal_conv1d_fn = None + + +class SSDBackend(str, Enum): + """ + Enum of the available state space scan implementations. + + Attributes: + NATIVE (str): The dependency-free pure-PyTorch chunked scan. Runs on CPU and GPU and is + ``torch.compile``-friendly, but noticeably slower and more memory hungry than the + fused kernels. + FUSED (str): The fused Triton kernels from ``mamba-ssm`` and ``causal-conv1d``. Requires + the optional ``mamba`` extra and a CUDA device. + """ + + NATIVE = "native" + FUSED = "fused" + + +# Above this model dimension the native backend becomes the throughput bottleneck and the fused +# kernels should be used instead. Chosen to flag real training runs without noise in unit tests. +_NATIVE_BACKEND_WARN_THRESHOLD_N_EMBD = 1024 + + +def _local_view(parameter: torch.Tensor) -> torch.Tensor: + """ + Returns the rank-local tensor backing a parameter. + + Under FSDP2 a parameter is a ``DTensor`` whose in-place ``copy_`` rejects a plain-tensor source. + Writing into the local shard side-steps that while leaving unsharded models unchanged. + + Args: + parameter (torch.Tensor): A parameter, possibly a ``DTensor``. + + Returns: + torch.Tensor: The local shard, or the parameter itself if it is not distributed. + """ + return parameter.to_local() if isinstance(parameter, DTensor) else parameter + + +class Mamba2Mixer(nn.Module): + """ + Mamba-2 selective state space mixer. + + A single input projection produces five packed tensors ``[z, x, B, C, dt]``. The ``x``, ``B`` + and ``C`` parts pass through a short causal depthwise convolution, the selective scan mixes + along the sequence, ``z`` gates the result via a grouped RMS norm, and an output projection + maps back to the model dimension. + """ + + def __init__( + self, + n_embd: int, + n_heads: int, + head_dim: int, + state_dim: int, + n_groups: int, + d_conv: int = 4, + chunk_size: int = 128, + ssd_backend: SSDBackend = SSDBackend.NATIVE, + norm_eps: float = 1e-5, + dt_min: float = 0.001, + dt_max: float = 0.1, + dt_init_floor: float = 1e-4, + A_init_range: tuple[float, float] = (1.0, 16.0), + bias: bool = False, + conv_bias: bool = True, + ): + """ + Initializes the Mamba2Mixer. + + Args: + n_embd (int): The model dimension. + n_heads (int): The number of SSM heads. Together with ``head_dim`` this determines the + inner dimension ``d_inner = n_heads * head_dim``. + head_dim (int): The dimension of a single SSM head. + state_dim (int): The SSM state dimension. + n_groups (int): The number of ``B``/``C`` groups. Must divide ``n_heads``. + d_conv (int): The kernel size of the causal depthwise convolution. + chunk_size (int): The chunk length of the selective scan. + ssd_backend (SSDBackend): Which scan implementation to use. + norm_eps (float): Epsilon of the gated RMS norm. + dt_min (float): Lower bound of the initial ``dt`` range. + dt_max (float): Upper bound of the initial ``dt`` range. + dt_init_floor (float): Lower clamp applied to the sampled ``dt`` before inverting + softplus. + A_init_range (tuple[float, float]): Range from which the per-head decay rates are + sampled uniformly before taking the logarithm. + bias (bool): Whether the input and output projections use a bias. + conv_bias (bool): Whether the depthwise convolution uses a bias. + + Raises: + ValueError: If the head/group configuration is inconsistent, or if the fused backend + is requested but ``mamba-ssm`` / ``causal-conv1d`` are not installed. + """ + super().__init__() + if n_heads % n_groups != 0: + raise ValueError(f"n_heads ({n_heads}) must be divisible by n_groups ({n_groups}).") + if d_conv < 1: + raise ValueError(f"d_conv must be at least 1, got {d_conv}.") + if A_init_range[0] <= 0 or A_init_range[1] < A_init_range[0]: + raise ValueError(f"A_init_range must satisfy 0 < low <= high, got {A_init_range}.") + + ssd_backend = SSDBackend(ssd_backend) + if ssd_backend == SSDBackend.FUSED: + missing = [ + name + for name, module in (("mamba-ssm", mamba_chunk_scan_combined), ("causal-conv1d", causal_conv1d_fn)) + if module is None + ] + if missing: + raise ValueError( + f"ssd_backend='fused' requires {' and '.join(missing)} to be installed. " + "Install the optional extra via `pip install -e '.[mamba]'`, or use " + "ssd_backend='native'." + ) + elif n_embd > _NATIVE_BACKEND_WARN_THRESHOLD_N_EMBD: + logger.warning( + "Mamba2Mixer is using the native (pure-PyTorch) SSD backend with n_embd=%d. " + "This is correct but significantly slower than the fused kernels. Consider " + "ssd_backend='fused' for large-scale training runs.", + n_embd, + ) + + self.n_embd = n_embd + self.n_heads = n_heads + self.head_dim = head_dim + self.state_dim = state_dim + self.n_groups = n_groups + self.d_conv = d_conv + self.chunk_size = chunk_size + self.ssd_backend = ssd_backend + self.d_inner = n_heads * head_dim + self.conv_dim = self.d_inner + 2 * n_groups * state_dim + self.dt_min = dt_min + self.dt_max = dt_max + self.dt_init_floor = dt_init_floor + self.A_init_range = A_init_range + + # in_proj packs [z, x, B, C, dt] into a single matmul. + self.in_proj = nn.Linear( + in_features=n_embd, + out_features=2 * self.d_inner + 2 * n_groups * state_dim + n_heads, + bias=bias, + ) + # The depthwise convolution is stored as raw parameters rather than an nn.Conv1d so that + # the fused and native paths can share the exact same layout. + self.conv1d_weight = nn.Parameter(torch.empty(self.conv_dim, 1, d_conv)) + self.conv1d_bias = nn.Parameter(torch.empty(self.conv_dim)) if conv_bias else None + # A_log and D are kept in fp32: they enter the scan through exp() and control its decay. + self.A_log = nn.Parameter(torch.empty(n_heads, dtype=torch.float32)) + self.D = nn.Parameter(torch.empty(n_heads, dtype=torch.float32)) + self.dt_bias = nn.Parameter(torch.empty(n_heads)) + self.norm = GatedRMSNorm(hidden_size=self.d_inner, num_groups=n_groups, eps=norm_eps) + self.out_proj = nn.Linear(in_features=self.d_inner, out_features=n_embd, bias=bias) + + self.reset_parameters() + + def reset_parameters(self) -> None: + """ + Initializes the SSM-specific parameters. + + The parameters ``A_log``, ``D``, ``dt_bias`` and the convolution kernel are *not* normally + distributed and must not be touched by the generic weight initialization. Modalities calls + ``reset_parameters`` on every submodule before running the configured initializer, and the + Nemotron parameter-name filters deliberately exclude these names, so this method is the + single place that defines their distribution. + + The distributions match the Megatron-LM reference implementation: + - ``A_log = log(U(A_init_range))`` + - ``dt_bias = softplus^-1(clamp(exp(U(log dt_min, log dt_max)), min=dt_init_floor))`` + - ``D = 1`` + - ``conv1d_weight`` / ``conv1d_bias``: the default nn.Conv1d initialization. + """ + if self.in_proj.weight.device.type == "meta": + # Nothing to initialize on the meta device; the model factory materializes the model + # and calls reset_parameters again afterwards. + return + + with torch.no_grad(): + # kaiming_uniform_/uniform_/ones_ are in-place ops that dispatch per shard and read the + # global shape, so they are safe on both plain tensors and FSDP2 DTensors. + nn.init.kaiming_uniform_(self.conv1d_weight, a=math.sqrt(5)) + if self.conv1d_bias is not None: + fan_in = self.conv1d_weight.size(1) * self.conv1d_weight.size(2) + bound = 1.0 / math.sqrt(fan_in) + nn.init.uniform_(self.conv1d_bias, -bound, bound) + nn.init.ones_(self.D) + + # dt_bias and A_log need values computed from a distribution that no nn.init helper + # provides, so they are written explicitly. Under FSDP2 the parameters are DTensors and + # copy_ rejects a plain-tensor source, so write into the local shard instead. The + # values are i.i.d. per head, so filling each rank's shard independently is equivalent + # to sampling the full tensor and slicing it. + dt_bias_local = _local_view(self.dt_bias) + A_log_local = _local_view(self.A_log) + + # Sample dt log-uniformly in [dt_min, dt_max], then invert softplus so that + # softplus(dt_bias) reproduces the sampled dt. + dt = torch.exp( + torch.rand(dt_bias_local.shape, device=dt_bias_local.device) + * (math.log(self.dt_max) - math.log(self.dt_min)) + + math.log(self.dt_min) + ).clamp(min=self.dt_init_floor) + dt_bias_local.copy_(dt + torch.log(-torch.expm1(-dt))) + + A = torch.empty(A_log_local.shape, device=A_log_local.device, dtype=torch.float32).uniform_( + *self.A_init_range + ) + A_log_local.copy_(torch.log(A).to(A_log_local.dtype)) + + def _split_projection( + self, zxbcdt: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Splits the packed input projection into its five components. + + Args: + zxbcdt (torch.Tensor): Packed projection of shape ``(B, L, 2 * d_inner + 2 * G * N + H)``. + + Returns: + tuple: ``(z, x, B, C, dt)`` with shapes ``(B, L, d_inner)``, ``(B, L, d_inner)``, + ``(B, L, G, N)``, ``(B, L, G, N)`` and ``(B, L, H)``. + """ + group_state_size = self.n_groups * self.state_dim + z, xbc, dt = torch.split( + zxbcdt, + [self.d_inner, self.d_inner + 2 * group_state_size, self.n_heads], + dim=-1, + ) + # The convolution operates on the concatenated [x, B, C] block. + xbc = self._apply_conv(xbc) + x, b, c = torch.split(xbc, [self.d_inner, group_state_size, group_state_size], dim=-1) + batch_size, seq_len = x.shape[:2] + b = b.view(batch_size, seq_len, self.n_groups, self.state_dim) + c = c.view(batch_size, seq_len, self.n_groups, self.state_dim) + return z, x, b, c, dt + + def _apply_conv(self, xbc: torch.Tensor) -> torch.Tensor: + """ + Applies the causal depthwise convolution followed by a SiLU activation. + + Args: + xbc (torch.Tensor): The concatenated ``[x, B, C]`` block of shape ``(B, L, conv_dim)``. + + Returns: + torch.Tensor: The activated convolution output of the same shape. + """ + if self.ssd_backend == SSDBackend.FUSED and causal_conv1d_fn is not None and xbc.is_cuda: + # causal_conv1d_fn expects (B, conv_dim, L) and folds the activation in. + y = causal_conv1d_fn( + x=xbc.transpose(1, 2).contiguous(), + weight=self.conv1d_weight.squeeze(1), + bias=self.conv1d_bias, + activation="silu", + ) + return y.transpose(1, 2) + return F.silu(causal_depthwise_conv1d(xbc, weight=self.conv1d_weight, bias=self.conv1d_bias)) + + def _run_scan( + self, + x: torch.Tensor, + dt: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + ) -> torch.Tensor: + """ + Runs the selective state space scan with the configured backend. + + Args: + x (torch.Tensor): SSM input of shape ``(B, L, H, P)``. + dt (torch.Tensor): Positive discretization step of shape ``(B, L, H)``. + b (torch.Tensor): Input projection of shape ``(B, L, G, N)``. + c (torch.Tensor): Output projection of shape ``(B, L, G, N)``. + + Returns: + torch.Tensor: The scan output of shape ``(B, L, H, P)``. + """ + A = -torch.exp(self.A_log.float()) + if self.ssd_backend == SSDBackend.FUSED and mamba_chunk_scan_combined is not None and x.is_cuda: + return mamba_chunk_scan_combined( + x, + dt, + A, + b, + c, + chunk_size=self.chunk_size, + D=self.D.float(), + z=None, + dt_bias=None, + dt_softplus=False, + ) + return ssd_chunked_scan( + x=x, + dt=dt, + A=A, + B=b, + C=c, + D=self.D, + chunk_size=self.chunk_size, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass of the Mamba-2 mixer. + + Args: + x (torch.Tensor): Input of shape ``(B, L, n_embd)``. + + Returns: + torch.Tensor: Output of shape ``(B, L, n_embd)``. + """ + batch_size, seq_len, _ = x.shape + z, x_ssm, b, c, dt = self._split_projection(self.in_proj(x)) + + # softplus keeps dt strictly positive; dt_bias sets its initial magnitude per head. + dt = F.softplus(dt.float() + self.dt_bias.float()) + x_heads = x_ssm.view(batch_size, seq_len, self.n_heads, self.head_dim) + + y = self._run_scan(x=x_heads, dt=dt, b=b, c=c) + y = y.reshape(batch_size, seq_len, self.d_inner) + y = self.norm(y, gate=z) + return self.out_proj(y) diff --git a/src/modalities/models/components/mamba2/ssd.py b/src/modalities/models/components/mamba2/ssd.py new file mode 100644 index 000000000..603d2c4ec --- /dev/null +++ b/src/modalities/models/components/mamba2/ssd.py @@ -0,0 +1,335 @@ +"""Pure-PyTorch Mamba-2 primitives (state space dual / SSD). + +This module provides a dependency-free, CPU-runnable and ``torch.compile``-friendly +implementation of the three building blocks a Mamba-2 mixer needs: + +1. :func:`ssd_chunked_scan` - the chunk-parallel selective state space scan. +2. :func:`causal_depthwise_conv1d` - the short causal depthwise convolution. +3. :class:`GatedRMSNorm` - the grouped, gated RMS normalization applied to the SSM output. + +The chunked scan follows the block-decomposition of the Mamba-2 paper (Dao & Gu, 2024) and is +numerically equivalent (up to floating point accumulation order) to the fused Triton kernel +``mamba_split_conv1d_scan_combined`` used by the Megatron-LM reference implementation. It is the +default backend so that Modalities does not take a hard dependency on ``mamba-ssm``; see +:class:`modalities.models.components.mamba2.mamba2_mixer.SSDBackend` for the fused alternative. + +Notation used throughout: + B: batch size, L: sequence length, H: number of SSM heads, P: SSM head dimension, + G: number of B/C groups, N: SSM state dimension, C: number of chunks, S: chunk length. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def _segment_sum(x: torch.Tensor) -> torch.Tensor: + """ + Computes the pairwise segment sums of a sequence of log-decay values. + + For an input of shape ``(..., T)`` this returns a tensor of shape ``(..., T, T)`` where + entry ``[..., i, j]`` holds ``sum(x[..., j + 1 : i + 1])`` for ``i >= j`` and ``-inf`` + otherwise. Exponentiating the result yields the causal decay matrix of the SSM, and doing + the accumulation in log space is what keeps long chunks numerically stable. + + Args: + x (torch.Tensor): Log-decay values of shape ``(..., T)``. + + Returns: + torch.Tensor: Segment sums of shape ``(..., T, T)``. + """ + seq_len = x.size(-1) + # Replicate along a new trailing axis so that entry [..., i, j] starts out as x[..., i]. + x_expanded = x.unsqueeze(-1).expand(*x.shape, seq_len) + strictly_lower = torch.tril(torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool), diagonal=-1) + x_expanded = x_expanded.masked_fill(~strictly_lower, 0) + # Accumulating over i turns entry [..., i, j] into sum(x[..., j + 1 : i + 1]). + segment_sums = torch.cumsum(x_expanded, dim=-2) + lower = torch.tril(torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool), diagonal=0) + return segment_sums.masked_fill(~lower, -torch.inf) + + +def _expand_groups_to_heads(x: torch.Tensor, num_heads: int) -> torch.Tensor: + """ + Expands a per-group tensor to a per-head tensor. + + Mamba-2 shares the ``B`` and ``C`` projections across all heads of a group (the SSM analogue + of grouped-query attention). Every group covers ``num_heads // num_groups`` consecutive heads. + + Args: + x (torch.Tensor): Per-group tensor of shape ``(B, L, G, N)``. + num_heads (int): The number of heads to expand to. + + Raises: + ValueError: If the number of heads is not divisible by the number of groups. + + Returns: + torch.Tensor: Per-head tensor of shape ``(B, L, H, N)``. + """ + num_groups = x.size(2) + if num_heads % num_groups != 0: + raise ValueError(f"num_heads ({num_heads}) must be divisible by num_groups ({num_groups}).") + heads_per_group = num_heads // num_groups + if heads_per_group == 1: + return x + return x.repeat_interleave(heads_per_group, dim=2) + + +def ssd_chunked_scan( + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + chunk_size: int = 128, + initial_states: torch.Tensor | None = None, + return_final_state: bool = False, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """ + Runs the chunk-parallel Mamba-2 selective state space scan. + + Implements the recurrence, for head ``h`` with associated group ``g``:: + + state_t = exp(dt_t * A_h) * state_{t-1} + dt_t * outer(B_t^g, x_t) + y_t = C_t^g @ state_t + D_h * x_t + + The sequence is split into chunks of ``chunk_size``. Within a chunk the recurrence is + evaluated as a masked matrix multiplication (the "quadratic"/attention-like form), while + across chunks only the compact ``(P, N)`` chunk states are passed. That is the duality the + Mamba-2 paper exploits, and it is what makes the scan efficient without a custom kernel. + + The scan is internally computed in float32 regardless of the input dtype, because the + ``cumsum`` over log-decays and the state accumulation are the numerically sensitive parts. + The result is cast back to the dtype of ``x``. + + Args: + x (torch.Tensor): SSM input of shape ``(B, L, H, P)``. + dt (torch.Tensor): Positive discretization step of shape ``(B, L, H)``, i.e. the value + after applying softplus to the projected ``dt`` plus its bias. + A (torch.Tensor): Negative per-head decay rate of shape ``(H,)``, i.e. ``-exp(A_log)``. + B (torch.Tensor): Input projection of shape ``(B, L, G, N)``. + C (torch.Tensor): Output projection of shape ``(B, L, G, N)``. + D (torch.Tensor | None): Optional per-head skip connection of shape ``(H,)``. + chunk_size (int): Chunk length used by the block decomposition. Must be positive. The + result is mathematically independent of this value; it only trades off compute + against memory. + initial_states (torch.Tensor | None): Optional initial SSM state of shape + ``(B, H, P, N)``. Defaults to zeros. + return_final_state (bool): Whether to additionally return the SSM state after the last + token, of shape ``(B, H, P, N)``. + + Raises: + ValueError: If ``chunk_size`` is not positive or the input shapes are inconsistent. + + Returns: + torch.Tensor | tuple[torch.Tensor, torch.Tensor]: The output of shape ``(B, L, H, P)``, + and the final state if ``return_final_state`` is set. + """ + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}.") + if x.dim() != 4: + raise ValueError(f"x must have shape (B, L, H, P), got {tuple(x.shape)}.") + if B.shape != C.shape: + raise ValueError(f"B and C must have the same shape, got {tuple(B.shape)} and {tuple(C.shape)}.") + + batch_size, seq_len, num_heads, head_dim = x.shape + out_dtype = x.dtype + + # The scan is accumulation-heavy; run it in fp32 and cast back at the end. + x_f = x.float() + dt_f = dt.float() + A_f = A.float() + B_f = _expand_groups_to_heads(B.float(), num_heads=num_heads) + C_f = _expand_groups_to_heads(C.float(), num_heads=num_heads) + + # Zero-pad the sequence up to a multiple of chunk_size. Appending zeros at the end is safe + # for a causal scan: padded tokens cannot influence the outputs of real tokens, and their + # own outputs are discarded. dt = 0 additionally makes the padded steps state-preserving. + pad_len = (-seq_len) % chunk_size + if pad_len > 0: + x_f = F.pad(x_f, (0, 0, 0, 0, 0, pad_len)) + dt_f = F.pad(dt_f, (0, 0, 0, pad_len)) + B_f = F.pad(B_f, (0, 0, 0, 0, 0, pad_len)) + C_f = F.pad(C_f, (0, 0, 0, 0, 0, pad_len)) + padded_len = seq_len + pad_len + num_chunks = padded_len // chunk_size + + # Fold dt into the input and the decay rate, matching the discretized formulation. + x_scaled = x_f * dt_f.unsqueeze(-1) # (B, L, H, P) + log_decay = A_f * dt_f # (B, L, H), non-positive + + # Reshape into chunks: (B, C, S, H, ...) + def to_chunks(tensor: torch.Tensor) -> torch.Tensor: + return tensor.reshape(batch_size, num_chunks, chunk_size, *tensor.shape[2:]) + + x_chunked = to_chunks(x_scaled) # (B, C, S, H, P) + B_chunked = to_chunks(B_f) # (B, C, S, H, N) + C_chunked = to_chunks(C_f) # (B, C, S, H, N) + # (B, C, S, H) -> (B, H, C, S) so that the decay cumsum runs over the chunk axis. + log_decay_chunked = to_chunks(log_decay).permute(0, 3, 1, 2) # (B, H, C, S) + log_decay_cumsum = torch.cumsum(log_decay_chunked, dim=-1) # (B, H, C, S) + + # 1. Intra-chunk (diagonal blocks): evaluate the recurrence in its quadratic form. + intra_chunk_decay = torch.exp(_segment_sum(log_decay_chunked)) # (B, H, C, S, S) + y_intra = torch.einsum( + "bclhn,bcshn,bhcls,bcshp->bclhp", C_chunked, B_chunked, intra_chunk_decay, x_chunked + ) # (B, C, S, H, P) + + # 2. Per-chunk end states: decay each position forward to the chunk boundary. + decay_to_chunk_end = torch.exp(log_decay_cumsum[..., -1:] - log_decay_cumsum) # (B, H, C, S) + chunk_states = torch.einsum("bcshn,bhcs,bcshp->bchpn", B_chunked, decay_to_chunk_end, x_chunked) # (B, C, H, P, N) + + # 3. Inter-chunk recurrence over the compact chunk states. + if initial_states is None: + leading_state = torch.zeros_like(chunk_states[:, :1]) # (B, 1, H, P, N) + else: + leading_state = initial_states.float().unsqueeze(1) # (B, 1, H, P, N) + all_states = torch.cat([leading_state, chunk_states], dim=1) # (B, C+1, H, P, N) + # Pad with a leading zero so that chunk 0 receives the (undecayed) initial state. + chunk_boundary_decay = torch.exp(_segment_sum(F.pad(log_decay_cumsum[..., -1], (1, 0)))) # (B, H, C+1, C+1) + propagated_states = torch.einsum("bhzc,bchpn->bzhpn", chunk_boundary_decay, all_states) # (B, C+1, H, P, N) + chunk_start_states = propagated_states[:, :-1] # (B, C, H, P, N) + + # 4. Inter-chunk (off-diagonal blocks): read the incoming chunk state out at every position. + decay_from_chunk_start = torch.exp(log_decay_cumsum) # (B, H, C, S) + y_inter = torch.einsum( + "bclhn,bchpn,bhcl->bclhp", C_chunked, chunk_start_states, decay_from_chunk_start + ) # (B, C, S, H, P) + + y = (y_intra + y_inter).reshape(batch_size, padded_len, num_heads, head_dim) + if pad_len > 0: + y = y[:, :seq_len] + + if D is not None: + y = y + D.float().view(1, 1, num_heads, 1) * x_f[:, :seq_len] + + y = y.to(out_dtype) + if return_final_state: + return y, propagated_states[:, -1].to(out_dtype) + return y + + +def ssd_recurrent_reference( + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, +) -> torch.Tensor: + """ + Evaluates the Mamba-2 recurrence step by step, as a correctness reference. + + This is the literal transcription of the SSM recurrence and is intentionally slow. It exists + so that :func:`ssd_chunked_scan` can be validated against a definition that involves no + block decomposition, no chunking and no log-space tricks. + + Args: + x (torch.Tensor): SSM input of shape ``(B, L, H, P)``. + dt (torch.Tensor): Positive discretization step of shape ``(B, L, H)``. + A (torch.Tensor): Negative per-head decay rate of shape ``(H,)``. + B (torch.Tensor): Input projection of shape ``(B, L, G, N)``. + C (torch.Tensor): Output projection of shape ``(B, L, G, N)``. + D (torch.Tensor | None): Optional per-head skip connection of shape ``(H,)``. + + Returns: + torch.Tensor: The output of shape ``(B, L, H, P)``. + """ + batch_size, seq_len, num_heads, head_dim = x.shape + x_f, dt_f, A_f = x.float(), dt.float(), A.float() + B_f = _expand_groups_to_heads(B.float(), num_heads=num_heads) + C_f = _expand_groups_to_heads(C.float(), num_heads=num_heads) + state_dim = B_f.size(-1) + + state = torch.zeros(batch_size, num_heads, head_dim, state_dim, device=x.device, dtype=torch.float32) + outputs = [] + for t in range(seq_len): + decay = torch.exp(dt_f[:, t] * A_f).view(batch_size, num_heads, 1, 1) # (B, H, 1, 1) + update = torch.einsum("bhp,bhn->bhpn", x_f[:, t] * dt_f[:, t].unsqueeze(-1), B_f[:, t]) + state = decay * state + update + outputs.append(torch.einsum("bhpn,bhn->bhp", state, C_f[:, t])) + y = torch.stack(outputs, dim=1) # (B, L, H, P) + + if D is not None: + y = y + D.float().view(1, 1, num_heads, 1) * x_f + return y.to(x.dtype) + + +def causal_depthwise_conv1d(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor: + """ + Applies a causal depthwise 1D convolution along the sequence dimension. + + Causality is achieved by left-padding with ``kernel_size - 1`` zeros so that position ``t`` + only ever sees positions ``<= t``. + + Args: + x (torch.Tensor): Input of shape ``(B, L, channels)``. + weight (torch.Tensor): Depthwise kernel of shape ``(channels, 1, kernel_size)``. + bias (torch.Tensor | None): Optional bias of shape ``(channels,)``. + + Returns: + torch.Tensor: Output of shape ``(B, L, channels)``. + """ + kernel_size = weight.size(-1) + channels = weight.size(0) + # conv1d expects (B, channels, L). + x_channels_first = x.transpose(1, 2) + x_padded = F.pad(x_channels_first, (kernel_size - 1, 0)) + y = F.conv1d(x_padded, weight=weight, bias=bias, groups=channels) + return y.transpose(1, 2) + + +class GatedRMSNorm(nn.Module): + """ + Grouped RMS normalization with a SiLU gate, as used inside the Mamba-2 mixer. + + The gate is applied *before* normalization (``norm_before_gate=False`` in the reference + implementation), and the normalization statistics are computed per group rather than over the + full inner dimension. Grouping matters because the SSM inner dimension is organized into + ``num_groups`` blocks whose scales can differ substantially. + """ + + def __init__(self, hidden_size: int, num_groups: int = 1, eps: float = 1e-5): + """ + Initializes the GatedRMSNorm module. + + Args: + hidden_size (int): The size of the normalized dimension (the SSM inner dimension). + num_groups (int): The number of groups to normalize independently. Must divide + ``hidden_size``. + eps (float): Value added to the mean square for numerical stability. + + Raises: + ValueError: If ``hidden_size`` is not divisible by ``num_groups``. + """ + super().__init__() + if hidden_size % num_groups != 0: + raise ValueError(f"hidden_size ({hidden_size}) must be divisible by num_groups ({num_groups}).") + self.hidden_size = hidden_size + self.num_groups = num_groups + self.group_size = hidden_size // num_groups + self.eps = eps + self.weight = nn.Parameter(torch.ones(hidden_size)) + + def reset_parameters(self) -> None: + """Resets the learnable scale to one. Called by the Modalities model initialization.""" + torch.nn.init.ones_(self.weight) + + def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """ + Applies the gate and the grouped RMS normalization. + + Args: + x (torch.Tensor): Input of shape ``(..., hidden_size)``. + gate (torch.Tensor): Gate of the same shape as ``x``. + + Returns: + torch.Tensor: Normalized output of the same shape and dtype as ``x``. + """ + out_dtype = x.dtype + gated = x.float() * F.silu(gate.float()) + grouped = gated.reshape(*gated.shape[:-1], self.num_groups, self.group_size) + inv_rms = torch.rsqrt(grouped.pow(2).mean(dim=-1, keepdim=True) + self.eps) + normalized = (grouped * inv_rms).reshape(*gated.shape) + return (normalized * self.weight.float()).to(out_dtype) diff --git a/src/modalities/models/components/moe/__init__.py b/src/modalities/models/components/moe/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/modalities/models/components/moe/experts.py b/src/modalities/models/components/moe/experts.py new file mode 100644 index 000000000..7ca4fbdd6 --- /dev/null +++ b/src/modalities/models/components/moe/experts.py @@ -0,0 +1,174 @@ +"""Grouped expert feed-forward networks for mixture-of-experts layers. + +All experts of a layer are stored in two stacked weight tensors so that the whole layer can be +evaluated with two grouped matrix multiplications instead of one small matmul per expert. This is +the same layout TorchTitan uses, and it is what makes 128 experts practical. + +Nemotron-3 Nano's experts are *not* gated: a single up-projection followed by a squared ReLU and a +down-projection. That halves the number of expert matrices compared to a SwiGLU expert. +""" + +import math +from enum import Enum + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class ExpertsBackend(str, Enum): + """ + Enum of the available grouped expert matmul implementations. + + Attributes: + GROUPED_MM (str): ``torch._grouped_mm``, a single kernel over all experts. Requires a + CUDA device and bfloat16 inputs; falls back to ``LOOPED`` automatically otherwise. + LOOPED (str): A Python loop over experts. Slower, but runs anywhere and in any dtype, + which makes it the reference implementation for tests. + """ + + GROUPED_MM = "grouped_mm" + LOOPED = "looped" + + +def squared_relu(x: torch.Tensor) -> torch.Tensor: + """ + Applies the squared ReLU activation, ``max(x, 0) ** 2``. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The activated tensor. + """ + return torch.square(F.relu(x)) + + +class GroupedExperts(nn.Module): + """ + A stack of ``num_experts`` non-gated feed-forward experts evaluated with grouped matmuls. + + The forward pass expects tokens that are already sorted by expert assignment, together with + the number of tokens per expert, so that each expert's slice is a contiguous block. + """ + + def __init__( + self, + n_embd: int, + ffn_hidden: int, + num_experts: int, + backend: ExpertsBackend = ExpertsBackend.GROUPED_MM, + ): + """ + Initializes the GroupedExperts module. + + Args: + n_embd (int): The model dimension. + ffn_hidden (int): The per-expert hidden dimension. + num_experts (int): The number of experts. + backend (ExpertsBackend): Which grouped matmul implementation to use. + """ + super().__init__() + self.n_embd = n_embd + self.ffn_hidden = ffn_hidden + self.num_experts = num_experts + self.backend = ExpertsBackend(backend) + + if self.backend == ExpertsBackend.GROUPED_MM: + # torch._grouped_mm requires every non-unit stride to be a multiple of 16 bytes. At + # 2 bytes per element that means both matmul inner dimensions must be multiples of 8. + # Checking here turns a cryptic mid-forward RuntimeError into a config-time error. + unaligned = { + name: value for name, value in (("n_embd", n_embd), ("ffn_hidden", ffn_hidden)) if value % 8 != 0 + } + if unaligned: + raise ValueError( + f"ExpertsBackend.GROUPED_MM requires 16-byte aligned strides, so " + f"{' and '.join(f'{name}={value}' for name, value in unaligned.items())} must be " + f"a multiple of 8. Adjust the dimensions or use ExpertsBackend.LOOPED." + ) + + # Expert dimension first, so that FSDP2 shards along the expert axis by default. + self.w1 = nn.Parameter(torch.empty(num_experts, ffn_hidden, n_embd)) + self.w2 = nn.Parameter(torch.empty(num_experts, n_embd, ffn_hidden)) + self.reset_parameters() + + def reset_parameters(self) -> None: + """ + Initializes the expert weights with the default ``nn.Linear`` scheme. + + These are raw ``nn.Parameter`` tensors rather than ``nn.Linear`` modules, so nothing else + would initialize them. Modalities calls ``reset_parameters`` on every submodule before the + configured weight initializer runs, which means a misconfigured regex filter degrades to + a sane default instead of leaving uninitialized memory in the model. + """ + if self.w1.device.type == "meta": + # Nothing to initialize on the meta device; the factory materializes and re-runs this. + return + with torch.no_grad(): + for weight in (self.w1, self.w2): + # Match nn.Linear: kaiming_uniform_ with a=sqrt(5) over the fan-in of each expert. + fan_in = weight.shape[-1] + bound = math.sqrt(1.0 / fan_in) * math.sqrt(3.0) + nn.init.uniform_(weight, -bound, bound) + + def _use_grouped_mm(self, x: torch.Tensor) -> bool: + """ + Decides whether the fused grouped matmul can be used for this input. + + Args: + x (torch.Tensor): The (sorted) token tensor. + + Returns: + bool: True if ``torch._grouped_mm`` is available and applicable. + """ + return ( + self.backend == ExpertsBackend.GROUPED_MM + and hasattr(torch, "_grouped_mm") + and x.is_cuda + and x.dtype in (torch.bfloat16, torch.float16) + ) + + def forward(self, x_sorted: torch.Tensor, tokens_per_expert: torch.Tensor) -> torch.Tensor: + """ + Applies each expert to its contiguous slice of the sorted token tensor. + + Args: + x_sorted (torch.Tensor): Tokens sorted by expert, of shape ``(num_routed, n_embd)``. + tokens_per_expert (torch.Tensor): Token count per expert, of shape ``(num_experts,)``, + integer dtype. Must sum to ``num_routed``. + + Returns: + torch.Tensor: Expert outputs of shape ``(num_routed, n_embd)``. + """ + if self._use_grouped_mm(x_sorted): + offsets = torch.cumsum(tokens_per_expert, dim=0, dtype=torch.int32) + hidden = squared_relu(torch._grouped_mm(x_sorted, self.w1.transpose(-2, -1), offs=offsets)) + return torch._grouped_mm(hidden, self.w2.transpose(-2, -1), offs=offsets) + + return self._forward_looped(x_sorted, tokens_per_expert) + + def _forward_looped(self, x_sorted: torch.Tensor, tokens_per_expert: torch.Tensor) -> torch.Tensor: + """ + Reference implementation: loop over experts and apply each to its slice. + + Args: + x_sorted (torch.Tensor): Tokens sorted by expert, of shape ``(num_routed, n_embd)``. + tokens_per_expert (torch.Tensor): Token count per expert, of shape ``(num_experts,)``. + + Returns: + torch.Tensor: Expert outputs of shape ``(num_routed, n_embd)``. + """ + outputs = torch.empty_like(x_sorted) + # A single host sync here instead of one per expert. + counts = tokens_per_expert.tolist() + start = 0 + for expert_idx, count in enumerate(counts): + if count == 0: + continue + stop = start + count + chunk = x_sorted[start:stop] + hidden = squared_relu(chunk @ self.w1[expert_idx].transpose(0, 1)) + outputs[start:stop] = hidden @ self.w2[expert_idx].transpose(0, 1) + start = stop + return outputs diff --git a/src/modalities/models/components/moe/load_balancing.py b/src/modalities/models/components/moe/load_balancing.py new file mode 100644 index 000000000..109487a6e --- /dev/null +++ b/src/modalities/models/components/moe/load_balancing.py @@ -0,0 +1,199 @@ +"""Auxiliary-loss-free load balancing for mixture-of-experts layers. + +Implements the update rule of "Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts" +(https://arxiv.org/abs/2408.15664), which Nemotron-3 Nano uses as its primary balancing mechanism. + +The idea: every expert carries an additive bias that is applied only when *selecting* experts, not +when weighting their outputs. After each optimizer step the bias of an under-loaded expert is nudged +up and that of an over-loaded expert down, by a fixed step size. Because the update depends only on +the *sign* of the load deviation, it is robust to the exact token counts and adds no gradient term +to the loss - unlike the classic auxiliary loss, it does not fight the language modelling objective. + +Two details matter for correctness in a distributed training loop: + +* The update must happen once per **optimizer step**, not per micro-batch, and it must see the load + summed over all data-parallel ranks. A rank-local update would let ranks drift apart, and a + per-micro-batch update would over-correct by the gradient accumulation factor. +* The token counters are reset after each update so that the next step measures a fresh window. +""" + +import logging + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed.device_mesh import DeviceMesh +from torch.optim import Optimizer + +from modalities.models.components.moe.moe import MoE +from modalities.running_env.fsdp.device_mesh import ParallelismDegrees + +logger = logging.getLogger(__name__) + + +def get_moe_layers(model: nn.Module) -> list[MoE]: + """ + Collects every mixture-of-experts block inside a model. + + Args: + model (nn.Module): The model to search. May be wrapped (FSDP1/FSDP2) or a pipeline stage. + + Returns: + list[MoE]: The MoE blocks found, in module traversal order. + """ + return [module for module in model.modules() if isinstance(module, MoE)] + + +# Name under which the non-pipeline mesh dimensions are flattened into a single reduction group. +_REDUCTION_MESH_DIM_NAME = "moe_load_balancing" + + +def get_expert_load_reduction_group(device_mesh: DeviceMesh | None) -> dist.ProcessGroup | None: + """ + Resolves the process group over which expert token counts have to be summed. + + All mesh dimensions except pipeline parallelism are included: + + * Data-parallel and context-parallel ranks each see different tokens, so their counts must be + added to obtain the true expert load. + * Tensor-parallel ranks see the *same* tokens. Including them scales every expert's count by the + same factor, and since the update rule only looks at the sign of the deviation from the mean, + a uniform factor is harmless. Including them keeps the collective simple. + * Pipeline-parallel ranks hold *different layers*. A given MoE layer exists on exactly one + stage, so a collective across pipeline ranks would try to match tensors that do not exist on + the other stages. Pipeline ranks must therefore be excluded. + + Args: + device_mesh (DeviceMesh | None): The device mesh, or None when running without one. + + Returns: + dist.ProcessGroup | None: The group to reduce over, or None if no reduction is needed. + """ + if not dist.is_available() or not dist.is_initialized(): + return None + if device_mesh is None: + return dist.group.WORLD + + mesh_dim_names = tuple(device_mesh.mesh_dim_names or ()) + reduction_dims = tuple(name for name in mesh_dim_names if name != ParallelismDegrees.PP.value) + if not reduction_dims: + # Pipeline parallelism only: every rank owns its layers alone, nothing to reduce. + return None + if len(reduction_dims) == len(mesh_dim_names): + return dist.group.WORLD + + submesh = device_mesh[reduction_dims] + if len(reduction_dims) == 1: + return submesh.get_group() + try: + return submesh._flatten(_REDUCTION_MESH_DIM_NAME).get_group() + except RuntimeError: + # Already flattened under this name by an earlier call. + return device_mesh[_REDUCTION_MESH_DIM_NAME].get_group() + + +@torch.no_grad() +def update_expert_biases( + moe_layers: list[MoE], + update_rate: float, + process_group: dist.ProcessGroup | None = None, +) -> None: + """ + Applies one auxiliary-loss-free load balancing step to every MoE layer. + + For each layer the per-expert token counts are summed across the data-parallel group, compared + against the mean load, and the expert bias is moved by ``update_rate`` in the direction that + equalizes the load. Counters are reset afterwards. + + Args: + moe_layers (list[MoE]): The MoE blocks to update. + update_rate (float): The step size of the bias update. Nemotron-3 Nano uses 1e-3. + process_group (dist.ProcessGroup | None): The group across which token counts are summed. + None means no reduction (single process, or already-global counts). + """ + for moe in moe_layers: + router = moe.router + if router.expert_bias is None: + continue + + tokens_per_expert = router.tokens_per_expert + if process_group is not None: + dist.all_reduce(tokens_per_expert, op=dist.ReduceOp.SUM, group=process_group) + + if tokens_per_expert.sum() == 0: + # No tokens were routed since the last update (e.g. an evaluation-only interval). + continue + + # Positive where an expert is under-loaded, negative where it is over-loaded. + load_deviation = tokens_per_expert.mean() - tokens_per_expert + router.expert_bias += torch.sign(load_deviation) * update_rate + router.tokens_per_expert.zero_() + + +class MoEBalancing: + """Factory for the load-balancing optimizer hook.""" + + @staticmethod + def register_expert_bias_update_hook( + optimizer: Optimizer, + model: nn.Module, + expert_bias_update_rate: float, + device_mesh: DeviceMesh | None = None, + ) -> Optimizer: + """ + Attaches the expert bias update to an optimizer and returns the same optimizer. + + This is an optimizer *decorator*: it is registered as an ``optimizer`` component variant so + that a config can wrap a plain Adam/AdamW without any change to the training loop. Using an + optimizer step pre-hook is what pins the update to one per optimizer step, making it correct + under gradient accumulation. + + Args: + optimizer (Optimizer): The optimizer to decorate. + model (nn.Module): The model whose MoE layers should be balanced. + expert_bias_update_rate (float): The step size of the bias update. + device_mesh (DeviceMesh | None): Device mesh used to resolve the data-parallel group. + + Raises: + ValueError: If the update rate is not positive. + + Returns: + Optimizer: The same optimizer instance, with the hook registered. + """ + if expert_bias_update_rate <= 0: + raise ValueError(f"expert_bias_update_rate must be positive, got {expert_bias_update_rate}.") + + moe_layers = get_moe_layers(model) + if not moe_layers: + logger.warning( + "No MoE layers found in the model; the expert bias update hook will not do anything. " + "(This is expected for a pipeline stage that holds no MoE layers.)" + ) + return optimizer + + balanced_layers = [moe for moe in moe_layers if moe.router.expert_bias is not None] + if not balanced_layers: + logger.warning( + "Found %d MoE layers but none of them maintains an expert bias " + "(use_expert_bias=False); auxiliary-loss-free load balancing is disabled.", + len(moe_layers), + ) + return optimizer + + process_group = get_expert_load_reduction_group(device_mesh) + + def _hook(opt: Optimizer, args: tuple, kwargs: dict) -> None: + del opt, args, kwargs + update_expert_biases( + moe_layers=balanced_layers, + update_rate=expert_bias_update_rate, + process_group=process_group, + ) + + optimizer.register_step_pre_hook(_hook) + logger.info( + "Registered auxiliary-loss-free load balancing for %d MoE layers (update rate %g).", + len(balanced_layers), + expert_bias_update_rate, + ) + return optimizer diff --git a/src/modalities/models/components/moe/moe.py b/src/modalities/models/components/moe/moe.py new file mode 100644 index 000000000..3c4b1ab0a --- /dev/null +++ b/src/modalities/models/components/moe/moe.py @@ -0,0 +1,158 @@ +"""Mixture-of-experts feed-forward layer. + +The layer combines a :class:`TopKRouter`, a :class:`GroupedExperts` stack and an optional always-on +shared expert. It implements the two load-balancing mechanisms that Nemotron-3 Nano uses together: + +1. **Auxiliary-loss-free balancing** (https://arxiv.org/abs/2408.15664): the layer counts how many + tokens each expert received; an optimizer hook nudges a per-expert selection bias towards the + mean load. This is the primary mechanism and requires no gradient plumbing. +2. **The classic load-balancing loss** (Switch Transformer / Lepikhin et al.): a small + differentiable penalty, computed per sequence, that discourages degenerate routing. Exposed via + :attr:`last_aux_loss` for the training loss to pick up. +""" + +import torch +import torch.nn as nn + +from modalities.models.components.moe.experts import GroupedExperts +from modalities.models.components.moe.router import TopKRouter + + +class MoE(nn.Module): + """ + A sparse mixture-of-experts feed-forward block. + + Tokens are routed to their top-k experts, sorted by expert so that each expert operates on a + contiguous block, evaluated with grouped matmuls, and scattered back with their routing + weights. A shared expert, if configured, processes every token densely and its output is added. + """ + + def __init__( + self, + router: TopKRouter, + experts: GroupedExperts, + shared_experts: nn.Module | None = None, + aux_loss_coeff: float = 0.0, + ): + """ + Initializes the MoE layer. + + Args: + router (TopKRouter): The routing module. + experts (GroupedExperts): The routed expert stack. + shared_experts (nn.Module | None): An optional dense module applied to every token. + Nemotron-3 Nano uses two shared experts, realized as one MLP of twice the expert + hidden dimension. + aux_loss_coeff (float): Coefficient of the sequence-level load-balancing loss. Zero + disables the loss (and its computation) entirely. + + Raises: + ValueError: If ``aux_loss_coeff`` is negative or the router and experts disagree on + the number of experts. + """ + super().__init__() + if aux_loss_coeff < 0: + raise ValueError(f"aux_loss_coeff must be non-negative, got {aux_loss_coeff}.") + if router.num_experts != experts.num_experts: + raise ValueError( + f"Router and experts disagree on the number of experts: " + f"{router.num_experts} vs {experts.num_experts}." + ) + + self.router = router + self.experts = experts + self.shared_experts = shared_experts + self.aux_loss_coeff = aux_loss_coeff + # Overwritten on every forward pass. Deliberately not a buffer: it must stay in the + # autograd graph and must not be checkpointed. + self.last_aux_loss: torch.Tensor | None = None + + def _compute_aux_loss(self, scores: torch.Tensor, top_indices: torch.Tensor, batch_size: int) -> torch.Tensor: + """ + Computes the sequence-level load-balancing loss. + + Follows ``switch_load_balancing_loss_func`` of the reference implementation:: + + loss = E * sum_i (f_i * P_i) + + with ``f_i`` the fraction of routed slots that went to expert ``i`` and ``P_i`` the mean + router probability of expert ``i``. Computing it per sequence and averaging over the batch + (rather than pooling the whole batch) is the ``seq_aux_loss`` variant: it prevents one + sequence's routing from being balanced out by another's. + + Args: + scores (torch.Tensor): Dense router scores of shape ``(num_tokens, num_experts)``. + top_indices (torch.Tensor): Selected experts of shape ``(num_tokens, top_k)``. + batch_size (int): The number of sequences in the batch. + + Returns: + torch.Tensor: The scalar auxiliary loss, already multiplied by ``aux_loss_coeff``. + """ + num_experts = self.router.num_experts + top_k = self.router.top_k + # Normalize over all experts so that P_i sums to one per token, matching the reference. + probs = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) + probs = probs.view(batch_size, -1, num_experts) + + routing_map = torch.zeros_like(scores).scatter_(-1, top_indices, 1.0) + routing_map = routing_map.view(batch_size, -1, num_experts) + + tokens_per_seq = probs.shape[1] + # f_i: fraction of the sequence's routed slots that landed on expert i. + fraction_per_expert = routing_map.sum(dim=1) / (tokens_per_seq * top_k) + # P_i: mean router probability assigned to expert i within the sequence. + prob_per_expert = probs.mean(dim=1) + per_sequence_loss = num_experts * (fraction_per_expert * prob_per_expert).sum(dim=-1) + return self.aux_loss_coeff * per_sequence_loss.mean() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass of the MoE layer. + + Args: + x (torch.Tensor): Input of shape ``(B, L, n_embd)``. + + Returns: + torch.Tensor: Output of shape ``(B, L, n_embd)``. + """ + batch_size, seq_len, n_embd = x.shape + x_flat = x.reshape(-1, n_embd) + + top_weights, top_indices, scores = self.router(x_flat) + + self.last_aux_loss = ( + self._compute_aux_loss(scores=scores, top_indices=top_indices, batch_size=batch_size) + if self.aux_loss_coeff > 0 + else None + ) + + # Flatten the (token, slot) pairs and sort them by expert so that every expert's tokens + # form one contiguous block, which is what the grouped matmul requires. + flat_expert_indices = top_indices.reshape(-1) + sort_order = torch.argsort(flat_expert_indices, stable=True) + sorted_expert_indices = flat_expert_indices[sort_order] + # Integer division recovers the token each routed slot belongs to. + sorted_token_indices = sort_order // self.router.top_k + + tokens_per_expert = torch.bincount(sorted_expert_indices, minlength=self.router.num_experts) + + # Track expert load for the auxiliary-loss-free bias update. Under activation + # checkpointing the forward pass runs twice, so counts are inflated by a constant factor. + # The bias update only uses the sign of the deviation from the mean, so this is harmless. + with torch.no_grad(): + self.router.tokens_per_expert += tokens_per_expert.to(self.router.tokens_per_expert.dtype) + + x_sorted = x_flat[sorted_token_indices] + expert_out = self.experts(x_sorted, tokens_per_expert) + + # Weight each routed slot by its routing probability and accumulate per token. + sorted_weights = top_weights.reshape(-1)[sort_order].unsqueeze(-1) + expert_out = expert_out * sorted_weights.to(expert_out.dtype) + + out_flat = torch.zeros_like(x_flat) + out_flat.index_add_(0, sorted_token_indices, expert_out.to(out_flat.dtype)) + out = out_flat.view(batch_size, seq_len, n_embd) + + if self.shared_experts is not None: + out = out + self.shared_experts(x) + return out diff --git a/src/modalities/models/components/moe/moe_losses.py b/src/modalities/models/components/moe/moe_losses.py new file mode 100644 index 000000000..493bc368b --- /dev/null +++ b/src/modalities/models/components/moe/moe_losses.py @@ -0,0 +1,129 @@ +"""Loss components for mixture-of-experts models. + +Nemotron-3 Nano combines the language modelling objective with a small load-balancing penalty. The +penalty itself is computed inside each MoE layer (it needs the router scores) and surfaced through +the model output dict; the components here read it back out and combine it with the main loss. +""" + +from typing import Annotated + +import torch +from pydantic import BaseModel, Field, model_validator + +from modalities.batch import InferenceResultBatch +from modalities.config.pydantic_if_types import PydanticLossIFType +from modalities.loss_functions import Loss + + +class MoEAuxLoss(Loss): + """ + Reads the pre-computed mixture-of-experts auxiliary loss out of the model output. + + The value is produced by the MoE layers during the forward pass (already multiplied by the + per-layer coefficient) and summed by the model. This class only surfaces it as a + :class:`~modalities.loss_functions.Loss` so that it can be logged and combined like any other + loss term. + """ + + def __init__(self, prediction_key: str, tag: str = "MoEAuxLoss"): + """ + Initializes the MoEAuxLoss. + + Args: + prediction_key (str): Key under which the model stores the summed auxiliary loss. + tag (str): Human-readable tag used for logging. + """ + super().__init__(tag) + self.prediction_key = prediction_key + + def __call__(self, forward_batch: InferenceResultBatch) -> torch.Tensor: + """ + Returns the auxiliary loss recorded in the forward batch. + + Args: + forward_batch (InferenceResultBatch): The model output. + + Returns: + torch.Tensor: The scalar auxiliary loss. + """ + return forward_batch.get_predictions(self.prediction_key) + + +class WeightedSumLoss(Loss): + """ + A weighted sum of several losses. + + Used to add the MoE load-balancing penalty to the language modelling loss without changing the + training loop, which evaluates exactly one loss component. + """ + + def __init__(self, losses: list[Loss], weights: list[float], tag: str = "WeightedSumLoss"): + """ + Initializes the WeightedSumLoss. + + Args: + losses (list[Loss]): The losses to combine. + weights (list[float]): One weight per loss. + tag (str): Human-readable tag used for logging. + + Raises: + ValueError: If the number of losses and weights differ, or no loss is given. + """ + super().__init__(tag) + if len(losses) == 0: + raise ValueError("WeightedSumLoss requires at least one loss.") + if len(losses) != len(weights): + raise ValueError(f"Got {len(losses)} losses but {len(weights)} weights; they must match.") + self.losses = losses + self.weights = weights + + def __call__(self, forward_batch: InferenceResultBatch) -> torch.Tensor: + """ + Computes the weighted sum of the configured losses. + + Args: + forward_batch (InferenceResultBatch): The model output. + + Returns: + torch.Tensor: The combined scalar loss. + """ + total = None + for loss, weight in zip(self.losses, self.weights, strict=True): + term = weight * loss(forward_batch) + total = term if total is None else total + term + return total + + +class MoEAuxLossConfig(BaseModel): + """ + Configuration of :class:`MoEAuxLoss`. + + Attributes: + prediction_key (str): Key under which the model stores the summed auxiliary loss. Must match + the model's ``aux_loss_key``. + tag (str): Human-readable tag used for logging. + """ + + prediction_key: str + tag: str = "MoEAuxLoss" + + +class WeightedSumLossConfig(BaseModel): + """ + Configuration of :class:`WeightedSumLoss`. + + Attributes: + losses (list[Loss]): The losses to combine. + weights (list[float]): One weight per loss. + tag (str): Human-readable tag used for logging. + """ + + losses: list[PydanticLossIFType] + weights: list[Annotated[float, Field(strict=True)]] + tag: str = "WeightedSumLoss" + + @model_validator(mode="after") + def _validate_lengths(self) -> "WeightedSumLossConfig": + if len(self.losses) != len(self.weights): + raise ValueError(f"Got {len(self.losses)} losses but {len(self.weights)} weights; they must match.") + return self diff --git a/src/modalities/models/components/moe/router.py b/src/modalities/models/components/moe/router.py new file mode 100644 index 000000000..10ad8c037 --- /dev/null +++ b/src/modalities/models/components/moe/router.py @@ -0,0 +1,130 @@ +"""Top-k mixture-of-experts router. + +Implements the routing scheme used by Nemotron-3 Nano: sigmoid gating with an additive +load-balancing bias that influences expert *selection* only, followed by a renormalization over +the selected experts and a constant rescaling. This matches the ``sigmoid`` branch of +``megatron.core.transformer.moe.moe_utils.topk_routing_with_score_function``. +""" + +from enum import Enum + +import torch +import torch.nn as nn + + +class RouterScoreFunction(str, Enum): + """ + Enum of the supported router score functions. + + Attributes: + SIGMOID (str): Independent per-expert sigmoid scores, renormalized over the selected + experts. Used by Nemotron-3 Nano and DeepSeek-style MoEs. + SOFTMAX (str): A softmax over the selected experts' logits. + """ + + SIGMOID = "sigmoid" + SOFTMAX = "softmax" + + +class TopKRouter(nn.Module): + """ + Routes each token to its top-k experts. + + The router is a single bias-free linear layer ("learnt MLP router"). Its scores are computed + in float32 regardless of the activation dtype, because top-k selection over 128 nearly-tied + scores is sensitive to rounding and any instability there shows up as expert flapping. + """ + + def __init__( + self, + n_embd: int, + num_experts: int, + top_k: int, + score_function: RouterScoreFunction = RouterScoreFunction.SIGMOID, + route_scale: float = 1.0, + use_expert_bias: bool = True, + router_dtype: torch.dtype = torch.float32, + ): + """ + Initializes the TopKRouter. + + Args: + n_embd (int): The model dimension. + num_experts (int): The number of routable experts. + top_k (int): How many experts each token is routed to. + score_function (RouterScoreFunction): Which score function to use. + route_scale (float): Constant factor applied to the final routing weights. Nemotron-3 + Nano uses 2.5 to compensate for the renormalization shrinking the weights. + use_expert_bias (bool): Whether to maintain an additive per-expert selection bias for + auxiliary-loss-free load balancing (https://arxiv.org/abs/2408.15664). + router_dtype (torch.dtype): Dtype in which scores and top-k selection are computed. + + Raises: + ValueError: If ``top_k`` is not in ``[1, num_experts]``. + """ + super().__init__() + if not 1 <= top_k <= num_experts: + raise ValueError(f"top_k must be in [1, num_experts={num_experts}], got {top_k}.") + + self.n_embd = n_embd + self.num_experts = num_experts + self.top_k = top_k + self.score_function = RouterScoreFunction(score_function) + self.route_scale = route_scale + self.router_dtype = router_dtype + + self.gate = nn.Linear(in_features=n_embd, out_features=num_experts, bias=False) + + if use_expert_bias: + # Persistent: the bias is part of the model state and must survive checkpointing. + self.register_buffer("expert_bias", torch.zeros(num_experts, dtype=torch.float32), persistent=True) + else: + self.expert_bias = None + # Non-persistent: a per-step counter that is reduced and reset by the load-balancing hook. + self.register_buffer("tokens_per_expert", torch.zeros(num_experts, dtype=torch.float32), persistent=False) + + def reset_parameters(self) -> None: + """Zeroes the load-balancing state. The gate weight is initialized by the model initializer.""" + with torch.no_grad(): + if self.expert_bias is not None and self.expert_bias.device.type != "meta": + self.expert_bias.zero_() + if self.tokens_per_expert.device.type != "meta": + self.tokens_per_expert.zero_() + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Computes the routing decision for every token. + + Args: + x (torch.Tensor): Input of shape ``(num_tokens, n_embd)``. + + Returns: + tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + - ``top_weights`` of shape ``(num_tokens, top_k)``: the routing weights, in the + dtype of ``x``. + - ``top_indices`` of shape ``(num_tokens, top_k)``: the selected expert indices. + - ``scores`` of shape ``(num_tokens, num_experts)``: the dense scores, kept for + the auxiliary load-balancing loss. + """ + logits = self.gate(x).to(self.router_dtype) + + if self.score_function == RouterScoreFunction.SIGMOID: + scores = torch.sigmoid(logits) + # The bias steers *selection* only; the returned weights come from the unbiased + # scores. Otherwise the balancing bias would leak into the forward signal. + selection_scores = scores if self.expert_bias is None else scores + self.expert_bias.to(scores.dtype) + _, top_indices = torch.topk(selection_scores, k=self.top_k, dim=-1) + top_weights = torch.gather(scores, dim=-1, index=top_indices) + if self.top_k > 1: + top_weights = top_weights / (top_weights.sum(dim=-1, keepdim=True) + 1e-20) + else: + scores = torch.softmax(logits, dim=-1) + selection_scores = scores if self.expert_bias is None else scores + self.expert_bias.to(scores.dtype) + _, top_indices = torch.topk(selection_scores, k=self.top_k, dim=-1) + top_logits = torch.gather(logits, dim=-1, index=top_indices) + top_weights = torch.softmax(top_logits, dim=-1) + + if self.route_scale != 1.0: + top_weights = top_weights * self.route_scale + + return top_weights.to(x.dtype), top_indices, scores diff --git a/src/modalities/models/components/norms.py b/src/modalities/models/components/norms.py new file mode 100644 index 000000000..477d00ce6 --- /dev/null +++ b/src/modalities/models/components/norms.py @@ -0,0 +1,58 @@ +"""Normalization wrapper config shared by models outside of GPT2. + +GPT2 carries its own, historically grown ``LayerNormWrapperConfig``. This module provides the +same declarative "which norm plus its config" indirection for newer models without coupling +them to the GPT2 implementation. +""" + +import torch.nn as nn +from pydantic import BaseModel + +from modalities.config.lookup_enum import LookupEnum +from modalities.models.components.layer_norms import ( + LayerNormConfig, + PytorchRMSLayerNormConfig, + RMSLayerNorm, + RMSLayerNormConfig, +) + + +class NormTypes(LookupEnum): + """ + Enum lookup class for normalization layers. + + Attributes: + rms_norm: The deprecated in-house RMSLayerNorm class. + layer_norm: nn.LayerNorm class. + pytorch_rms_norm: nn.RMSNorm class. + """ + + rms_norm = RMSLayerNorm + layer_norm = nn.LayerNorm + pytorch_rms_norm = nn.RMSNorm + + +class NormWrapperConfig(BaseModel): + """ + Configuration selecting a normalization layer type together with its own configuration. + + Attributes: + norm_type (NormTypes): Which normalization implementation to instantiate. + config (PytorchRMSLayerNormConfig | RMSLayerNormConfig | LayerNormConfig): The + constructor arguments of the selected normalization implementation. + """ + + norm_type: NormTypes + config: PytorchRMSLayerNormConfig | RMSLayerNormConfig | LayerNormConfig + + def build(self) -> nn.Module: + """ + Instantiates a fresh normalization module. + + A new instance is created on every call so that each layer of a model owns its own + normalization parameters. + + Returns: + nn.Module: The instantiated normalization module. + """ + return self.norm_type.value(**dict(self.config)) diff --git a/src/modalities/models/nemotron/__init__.py b/src/modalities/models/nemotron/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/modalities/models/nemotron/layer_pattern.py b/src/modalities/models/nemotron/layer_pattern.py new file mode 100644 index 000000000..a9f1b543d --- /dev/null +++ b/src/modalities/models/nemotron/layer_pattern.py @@ -0,0 +1,90 @@ +"""Layer pattern parsing for hybrid Mamba-Transformer models. + +A hybrid model such as Nemotron-3 Nano is described by a pattern string in which every +character denotes one residual sublayer, e.g. ``"MEM*E"``. Each sublayer is a self-contained +pre-norm residual block (``x = x + f(norm(x))``); there is no combined "attention + MLP" +block as in a classical transformer. + +The symbols follow the reference implementation in Megatron-LM +(``megatron/core/models/hybrid/hybrid_layer_allocation.py``) so that pattern strings can be +copied verbatim between the two frameworks. +""" + +from enum import Enum + + +class LayerSymbol(str, Enum): + """ + Enum of the layer types that a hybrid layer pattern can contain. + + The values match the symbols used by Megatron-LM so that pattern strings are portable. + + Attributes: + MAMBA (str): A Mamba-2 mixer layer. + ATTENTION (str): A (grouped-query) self-attention layer. + MOE (str): A mixture-of-experts feed-forward layer. + MLP (str): A dense feed-forward layer. + """ + + MAMBA = "M" + ATTENTION = "*" + MOE = "E" + MLP = "-" + + +def parse_layer_pattern(pattern: str) -> list[LayerSymbol]: + """ + Parses a hybrid layer pattern string into a list of layer symbols. + + Args: + pattern (str): The pattern string, e.g. ``"MEM*E"``. + + Raises: + ValueError: If the pattern is empty or contains an unknown symbol. + + Returns: + list[LayerSymbol]: One symbol per layer, in model order. + """ + if len(pattern) == 0: + raise ValueError("The layer pattern must not be empty.") + + valid_symbols = {symbol.value for symbol in LayerSymbol} + layer_symbols: list[LayerSymbol] = [] + for position, character in enumerate(pattern): + if character not in valid_symbols: + raise ValueError( + f"Invalid layer symbol '{character}' at position {position} of layer pattern '{pattern}'. " + f"Valid symbols are {sorted(valid_symbols)}." + ) + layer_symbols.append(LayerSymbol(character)) + return layer_symbols + + +def count_layers_by_type(pattern: str) -> dict[LayerSymbol, int]: + """ + Counts how many layers of each type a pattern contains. + + Args: + pattern (str): The pattern string, e.g. ``"MEM*E"``. + + Returns: + dict[LayerSymbol, int]: Counts for every layer type, including types with a count of zero. + """ + layer_symbols = parse_layer_pattern(pattern) + counts = {symbol: 0 for symbol in LayerSymbol} + for layer_symbol in layer_symbols: + counts[layer_symbol] += 1 + return counts + + +def get_num_layers(pattern: str) -> int: + """ + Returns the total number of layers described by a pattern. + + Args: + pattern (str): The pattern string, e.g. ``"MEM*E"``. + + Returns: + int: The number of layers. + """ + return len(parse_layer_pattern(pattern)) diff --git a/src/modalities/models/nemotron/nemotron_attention.py b/src/modalities/models/nemotron/nemotron_attention.py new file mode 100644 index 000000000..50518f069 --- /dev/null +++ b/src/modalities/models/nemotron/nemotron_attention.py @@ -0,0 +1,176 @@ +"""Grouped-query causal self-attention with a head dimension decoupled from the model dimension. + +This cannot reuse the GPT2 attention module: GPT2 derives the head dimension as +``n_embd // n_head_q``, whereas Nemotron-3 Nano uses 32 query heads of dimension 128 on a model +dimension of 2688 (32 * 128 = 4096, which is deliberately larger than 2688). Attention here also +applies no positional transform at all, since the Mamba layers supply positional information. +""" + +import math +from enum import Enum + +import torch +import torch.nn as nn + +try: + from flash_attn import flash_attn_func +except ModuleNotFoundError: + flash_attn_func = None + + +class NemotronAttentionImplementation(str, Enum): + """ + Enum of the supported attention kernels. + + Attributes: + MANUAL (str): An explicit softmax implementation. Slow, but runs anywhere and is used as + the reference in tests. + PYTORCH_FLASH (str): ``torch.nn.functional.scaled_dot_product_attention``. + DAO_FLASH (str): The ``flash-attn`` package's kernel. + """ + + MANUAL = "manual" + PYTORCH_FLASH = "pytorch_flash" + DAO_FLASH = "dao_flash" + + +class NemotronSelfAttention(nn.Module): + """ + Causal grouped-query self-attention. + + Query, key and value projections are separate and bias-free. The number of key/value heads may + be smaller than the number of query heads (grouped-query attention); key and value heads are + repeated to match the queries before the attention kernel is invoked. + """ + + def __init__( + self, + n_embd: int, + n_head_q: int, + n_head_kv: int, + head_dim: int, + attention_implementation: NemotronAttentionImplementation = NemotronAttentionImplementation.PYTORCH_FLASH, + bias: bool = False, + dropout: float = 0.0, + ): + """ + Initializes the NemotronSelfAttention module. + + Args: + n_embd (int): The model dimension. + n_head_q (int): The number of query heads. + n_head_kv (int): The number of key/value heads. Must divide ``n_head_q``. + head_dim (int): The dimension of a single attention head. Independent of ``n_embd``. + attention_implementation (NemotronAttentionImplementation): Which kernel to use. + bias (bool): Whether the projections use a bias. + dropout (float): Attention dropout probability. + + Raises: + ValueError: If ``n_head_q`` is not divisible by ``n_head_kv``. + """ + super().__init__() + if n_head_q % n_head_kv != 0: + raise ValueError(f"n_head_q ({n_head_q}) must be divisible by n_head_kv ({n_head_kv}).") + + self.n_embd = n_embd + self.n_head_q = n_head_q + self.n_head_kv = n_head_kv + self.head_dim = head_dim + self.n_rep = n_head_q // n_head_kv + self.dropout = dropout + self.attention_implementation = NemotronAttentionImplementation(attention_implementation) + + self.q_attn = nn.Linear(in_features=n_embd, out_features=n_head_q * head_dim, bias=bias) + self.k_attn = nn.Linear(in_features=n_embd, out_features=n_head_kv * head_dim, bias=bias) + self.v_attn = nn.Linear(in_features=n_embd, out_features=n_head_kv * head_dim, bias=bias) + self.c_proj = nn.Linear(in_features=n_head_q * head_dim, out_features=n_embd, bias=bias) + self.resid_dropout = nn.Dropout(dropout) + + @staticmethod + def _repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + Repeats each key/value head ``n_rep`` times to match the number of query heads. + + Args: + x (torch.Tensor): Tensor of shape ``(B, n_head_kv, T, head_dim)``. + n_rep (int): The repetition factor. + + Returns: + torch.Tensor: Tensor of shape ``(B, n_head_kv * n_rep, T, head_dim)``. + """ + if n_rep == 1: + return x + batch_size, n_head_kv, seq_len, head_dim = x.shape + return ( + x[:, :, None, :, :] + .expand(batch_size, n_head_kv, n_rep, seq_len, head_dim) + .reshape(batch_size, n_head_kv * n_rep, seq_len, head_dim) + ) + + def _execute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + """ + Runs the configured causal attention kernel. + + Args: + q (torch.Tensor): Queries of shape ``(B, n_head_q, T, head_dim)``. + k (torch.Tensor): Keys of shape ``(B, n_head_kv, T, head_dim)``. + v (torch.Tensor): Values of shape ``(B, n_head_kv, T, head_dim)``. + + Raises: + NotImplementedError: If the DAO flash attention kernel is requested but not installed. + + Returns: + torch.Tensor: Attention output of shape ``(B, T, n_head_q, head_dim)``. + """ + dropout_p = self.dropout if self.training else 0.0 + + if self.attention_implementation == NemotronAttentionImplementation.DAO_FLASH: + if flash_attn_func is None: + raise NotImplementedError("Dao flash attention is requested but flash-attn is not installed.") + # flash_attn_func handles grouped-query attention natively and wants (B, T, H, hd). + return flash_attn_func( + q.transpose(1, 2).contiguous(), + k.transpose(1, 2).contiguous(), + v.transpose(1, 2).contiguous(), + dropout_p=dropout_p, + causal=True, + ) + + k = self._repeat_kv(k, self.n_rep) + v = self._repeat_kv(v, self.n_rep) + + if self.attention_implementation == NemotronAttentionImplementation.PYTORCH_FLASH: + y = torch.nn.functional.scaled_dot_product_attention( + query=q, key=k, value=v, attn_mask=None, dropout_p=dropout_p, is_causal=True + ) + else: + seq_len = q.size(-2) + scale = 1.0 / math.sqrt(q.size(-1)) + attn_weights = (q @ k.transpose(-2, -1)) * scale + causal_mask = torch.ones(seq_len, seq_len, dtype=torch.bool, device=q.device).tril() + attn_weights = attn_weights.masked_fill(~causal_mask, float("-inf")) + attn_weights = torch.softmax(attn_weights, dim=-1) + attn_weights = torch.dropout(attn_weights, dropout_p, train=self.training) + y = attn_weights @ v + + return y.transpose(1, 2).contiguous() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass of the attention module. + + Args: + x (torch.Tensor): Input of shape ``(B, T, n_embd)``. + + Returns: + torch.Tensor: Output of shape ``(B, T, n_embd)``. + """ + batch_size, seq_len, _ = x.shape + + q = self.q_attn(x).view(batch_size, seq_len, self.n_head_q, self.head_dim).transpose(1, 2) + k = self.k_attn(x).view(batch_size, seq_len, self.n_head_kv, self.head_dim).transpose(1, 2) + v = self.v_attn(x).view(batch_size, seq_len, self.n_head_kv, self.head_dim).transpose(1, 2) + + y = self._execute_attention(q, k, v) + y = y.reshape(batch_size, seq_len, self.n_head_q * self.head_dim) + return self.resid_dropout(self.c_proj(y)) diff --git a/src/modalities/models/nemotron/nemotron_layer_specs.py b/src/modalities/models/nemotron/nemotron_layer_specs.py new file mode 100644 index 000000000..54fc91124 --- /dev/null +++ b/src/modalities/models/nemotron/nemotron_layer_specs.py @@ -0,0 +1,347 @@ +"""Layer specifications: configurable builders for the sublayers of a hybrid model. + +A hybrid model contains dozens of layers of a handful of types. Modalities' component factory +instantiates every config node exactly once, so injecting *instantiated* layer modules would make +all layers of a type share one set of weights. Layer specs solve this the same way Megatron-LM's +``ModuleSpec`` does: the registry produces a declarative builder, and the model calls +:meth:`NemotronLayerSpecIF.build` once per layer position to get a fresh module. + +Every network component (mixer, attention, experts, router, norms) is therefore fully configurable +from YAML while still yielding independent parameters per layer. +""" + +from abc import ABC, abstractmethod +from typing import Annotated, Optional + +import torch +import torch.nn as nn +from pydantic import BaseModel, Field, model_validator + +from modalities.models.components.mamba2.mamba2_mixer import Mamba2Mixer, SSDBackend +from modalities.models.components.moe.experts import ExpertsBackend, GroupedExperts +from modalities.models.components.moe.moe import MoE +from modalities.models.components.moe.router import RouterScoreFunction, TopKRouter +from modalities.models.components.norms import NormWrapperConfig +from modalities.models.nemotron.layer_pattern import LayerSymbol +from modalities.models.nemotron.nemotron_attention import NemotronAttentionImplementation, NemotronSelfAttention +from modalities.models.nemotron.nemotron_layers import ( + Mamba2Layer, + NemotronAttentionLayer, + NemotronMLPLayer, + NemotronMoELayer, +) +from modalities.models.nemotron.nemotron_mlp import SquaredReLUMLP + +PositiveInt = Annotated[int, Field(strict=True, ge=1)] + + +class NemotronLayerSpecIF(ABC): + """Interface for a builder that produces one residual sublayer per model position.""" + + @property + @abstractmethod + def symbol(self) -> LayerSymbol: + """ + The layer pattern symbol this spec is responsible for. + + Returns: + LayerSymbol: The symbol, e.g. ``LayerSymbol.MAMBA`` for ``"M"``. + """ + raise NotImplementedError + + @abstractmethod + def build(self, layer_idx: int) -> nn.Module: + """ + Builds a fresh layer module with its own parameters. + + Args: + layer_idx (int): The index of the layer within the model. Passed through for specs + that want depth-dependent behaviour; the current specs ignore it. + + Returns: + nn.Module: A newly constructed layer. + """ + raise NotImplementedError + + +class Mamba2LayerSpecConfig(BaseModel): + """ + Configuration of a Mamba-2 layer. + + Attributes: + n_embd (int): The model dimension. + mamba_n_heads (int): Number of SSM heads. + mamba_head_dim (int): Dimension of a single SSM head. + mamba_state_dim (int): SSM state dimension. + mamba_n_groups (int): Number of B/C groups; must divide ``mamba_n_heads``. + d_conv (int): Kernel size of the causal depthwise convolution. + chunk_size (int): Chunk length of the selective scan. + ssd_backend (SSDBackend): Which scan implementation to use. + norm_config (NormWrapperConfig): The pre-normalization of the layer. + mixer_norm_eps (float): Epsilon of the mixer-internal gated RMS norm. + bias (bool): Whether the input/output projections use a bias. + conv_bias (bool): Whether the depthwise convolution uses a bias. + """ + + n_embd: PositiveInt + mamba_n_heads: PositiveInt + mamba_head_dim: PositiveInt + mamba_state_dim: PositiveInt + mamba_n_groups: PositiveInt + d_conv: PositiveInt = 4 + chunk_size: PositiveInt = 128 + ssd_backend: SSDBackend = SSDBackend.NATIVE + norm_config: NormWrapperConfig + mixer_norm_eps: Annotated[float, Field(strict=True, gt=0)] = 1e-5 + bias: bool = False + conv_bias: bool = True + + @model_validator(mode="after") + def _validate_head_group_divisibility(self) -> "Mamba2LayerSpecConfig": + if self.mamba_n_heads % self.mamba_n_groups != 0: + raise ValueError( + f"mamba_n_heads ({self.mamba_n_heads}) must be divisible by " f"mamba_n_groups ({self.mamba_n_groups})." + ) + return self + + +class Mamba2LayerSpec(NemotronLayerSpecIF): + """Builds :class:`Mamba2Layer` instances from a :class:`Mamba2LayerSpecConfig`.""" + + def __init__(self, **config): + """ + Initializes the spec. + + Args: + **config: The fields of :class:`Mamba2LayerSpecConfig`. + """ + self.config = Mamba2LayerSpecConfig(**config) + + @property + def symbol(self) -> LayerSymbol: + return LayerSymbol.MAMBA + + def build(self, layer_idx: int) -> Mamba2Layer: + config = self.config + mixer = Mamba2Mixer( + n_embd=config.n_embd, + n_heads=config.mamba_n_heads, + head_dim=config.mamba_head_dim, + state_dim=config.mamba_state_dim, + n_groups=config.mamba_n_groups, + d_conv=config.d_conv, + chunk_size=config.chunk_size, + ssd_backend=config.ssd_backend, + norm_eps=config.mixer_norm_eps, + bias=config.bias, + conv_bias=config.conv_bias, + ) + return Mamba2Layer(norm=config.norm_config.build(), mixer=mixer) + + +class NemotronAttentionLayerSpecConfig(BaseModel): + """ + Configuration of a grouped-query self-attention layer. + + Attributes: + n_embd (int): The model dimension. + n_head_q (int): Number of query heads. + n_head_kv (int): Number of key/value heads; must divide ``n_head_q``. + head_dim (int): Dimension of a single attention head, independent of ``n_embd``. + attention_implementation (NemotronAttentionImplementation): Which kernel to use. + norm_config (NormWrapperConfig): The pre-normalization of the layer. + bias (bool): Whether the projections use a bias. + dropout (float): Attention dropout probability. + """ + + n_embd: PositiveInt + n_head_q: PositiveInt + n_head_kv: PositiveInt + head_dim: PositiveInt + attention_implementation: NemotronAttentionImplementation = NemotronAttentionImplementation.PYTORCH_FLASH + norm_config: NormWrapperConfig + bias: bool = False + dropout: Annotated[float, Field(strict=True, ge=0.0, lt=1.0)] = 0.0 + + @model_validator(mode="after") + def _validate_head_divisibility(self) -> "NemotronAttentionLayerSpecConfig": + if self.n_head_q % self.n_head_kv != 0: + raise ValueError(f"n_head_q ({self.n_head_q}) must be divisible by n_head_kv ({self.n_head_kv}).") + return self + + +class NemotronAttentionLayerSpec(NemotronLayerSpecIF): + """Builds :class:`NemotronAttentionLayer` instances.""" + + def __init__(self, **config): + """ + Initializes the spec. + + Args: + **config: The fields of :class:`NemotronAttentionLayerSpecConfig`. + """ + self.config = NemotronAttentionLayerSpecConfig(**config) + + @property + def symbol(self) -> LayerSymbol: + return LayerSymbol.ATTENTION + + def build(self, layer_idx: int) -> NemotronAttentionLayer: + config = self.config + attn = NemotronSelfAttention( + n_embd=config.n_embd, + n_head_q=config.n_head_q, + n_head_kv=config.n_head_kv, + head_dim=config.head_dim, + attention_implementation=config.attention_implementation, + bias=config.bias, + dropout=config.dropout, + ) + return NemotronAttentionLayer(norm=config.norm_config.build(), attn=attn) + + +class NemotronMoELayerSpecConfig(BaseModel): + """ + Configuration of a mixture-of-experts layer. + + Attributes: + n_embd (int): The model dimension. + num_experts (int): Number of routed experts. + moe_ffn_hidden (int): Hidden dimension of a single routed expert. + top_k (int): Number of experts each token is routed to. + score_function (RouterScoreFunction): Router score function. + route_scale (float): Constant factor applied to the routing weights. + use_expert_bias (bool): Whether to maintain the auxiliary-loss-free balancing bias. + router_dtype (str): Dtype for router score computation, ``"float32"`` or ``"bfloat16"``. + num_shared_experts (int): Number of always-on shared experts. Realized as a single MLP of + ``num_shared_experts * shared_expert_ffn_hidden_per_expert`` hidden units, matching the + reference implementation. Zero disables shared experts. + shared_expert_ffn_hidden_per_expert (int | None): Hidden dimension of one shared expert. + Defaults to ``moe_ffn_hidden``. + aux_loss_coeff (float): Coefficient of the sequence-level load-balancing loss. + experts_backend (ExpertsBackend): Which grouped matmul implementation to use. + norm_config (NormWrapperConfig): The pre-normalization of the layer. + bias (bool): Whether the shared expert projections use a bias. + """ + + n_embd: PositiveInt + num_experts: PositiveInt + moe_ffn_hidden: PositiveInt + top_k: PositiveInt + score_function: RouterScoreFunction = RouterScoreFunction.SIGMOID + route_scale: Annotated[float, Field(strict=True, gt=0.0)] = 1.0 + use_expert_bias: bool = True + router_dtype: str = "float32" + num_shared_experts: Annotated[int, Field(strict=True, ge=0)] = 0 + shared_expert_ffn_hidden_per_expert: Optional[PositiveInt] = None + aux_loss_coeff: Annotated[float, Field(strict=True, ge=0.0)] = 0.0 + experts_backend: ExpertsBackend = ExpertsBackend.GROUPED_MM + norm_config: NormWrapperConfig + bias: bool = False + + @model_validator(mode="after") + def _validate(self) -> "NemotronMoELayerSpecConfig": + if self.top_k > self.num_experts: + raise ValueError(f"top_k ({self.top_k}) must not exceed num_experts ({self.num_experts}).") + if self.router_dtype not in ("float32", "bfloat16"): + raise ValueError(f"router_dtype must be 'float32' or 'bfloat16', got '{self.router_dtype}'.") + return self + + @property + def shared_expert_ffn_hidden(self) -> int: + """ + The total hidden dimension of the fused shared expert MLP. + + Returns: + int: ``num_shared_experts * shared_expert_ffn_hidden_per_expert``, or 0 if disabled. + """ + if self.num_shared_experts == 0: + return 0 + per_expert = self.shared_expert_ffn_hidden_per_expert or self.moe_ffn_hidden + return self.num_shared_experts * per_expert + + +class NemotronMoELayerSpec(NemotronLayerSpecIF): + """Builds :class:`NemotronMoELayer` instances.""" + + def __init__(self, **config): + """ + Initializes the spec. + + Args: + **config: The fields of :class:`NemotronMoELayerSpecConfig`. + """ + self.config = NemotronMoELayerSpecConfig(**config) + + @property + def symbol(self) -> LayerSymbol: + return LayerSymbol.MOE + + def build(self, layer_idx: int) -> NemotronMoELayer: + config = self.config + router = TopKRouter( + n_embd=config.n_embd, + num_experts=config.num_experts, + top_k=config.top_k, + score_function=config.score_function, + route_scale=config.route_scale, + use_expert_bias=config.use_expert_bias, + router_dtype=getattr(torch, config.router_dtype), + ) + experts = GroupedExperts( + n_embd=config.n_embd, + ffn_hidden=config.moe_ffn_hidden, + num_experts=config.num_experts, + backend=config.experts_backend, + ) + shared_experts = ( + SquaredReLUMLP(n_embd=config.n_embd, ffn_hidden=config.shared_expert_ffn_hidden, bias=config.bias) + if config.num_shared_experts > 0 + else None + ) + moe = MoE( + router=router, + experts=experts, + shared_experts=shared_experts, + aux_loss_coeff=config.aux_loss_coeff, + ) + return NemotronMoELayer(norm=config.norm_config.build(), moe=moe) + + +class NemotronMLPLayerSpecConfig(BaseModel): + """ + Configuration of a dense squared-ReLU feed-forward layer. + + Attributes: + n_embd (int): The model dimension. + ffn_hidden (int): The hidden dimension. + norm_config (NormWrapperConfig): The pre-normalization of the layer. + bias (bool): Whether the projections use a bias. + """ + + n_embd: PositiveInt + ffn_hidden: PositiveInt + norm_config: NormWrapperConfig + bias: bool = False + + +class NemotronMLPLayerSpec(NemotronLayerSpecIF): + """Builds :class:`NemotronMLPLayer` instances.""" + + def __init__(self, **config): + """ + Initializes the spec. + + Args: + **config: The fields of :class:`NemotronMLPLayerSpecConfig`. + """ + self.config = NemotronMLPLayerSpecConfig(**config) + + @property + def symbol(self) -> LayerSymbol: + return LayerSymbol.MLP + + def build(self, layer_idx: int) -> NemotronMLPLayer: + config = self.config + mlp = SquaredReLUMLP(n_embd=config.n_embd, ffn_hidden=config.ffn_hidden, bias=config.bias) + return NemotronMLPLayer(norm=config.norm_config.build(), mlp=mlp) diff --git a/src/modalities/models/nemotron/nemotron_layers.py b/src/modalities/models/nemotron/nemotron_layers.py new file mode 100644 index 000000000..76830497a --- /dev/null +++ b/src/modalities/models/nemotron/nemotron_layers.py @@ -0,0 +1,129 @@ +"""The four residual sublayer types of a hybrid Mamba-Transformer stack. + +Unlike a classical transformer block, which bundles attention and a feed-forward network, every +layer here wraps exactly *one* operator in a pre-norm residual connection:: + + x = x + operator(norm(x)) + +A model is then a sequence of such single-operator layers described by a layer pattern, e.g. +``"MEM*E"``. This is the structure of Nemotron-H and Nemotron-3 Nano, and it is what allows the +Mamba / attention / MoE ratio to be tuned independently. +""" + +import torch +import torch.nn as nn + +from modalities.models.components.mamba2.mamba2_mixer import Mamba2Mixer +from modalities.models.components.moe.moe import MoE +from modalities.models.nemotron.nemotron_attention import NemotronSelfAttention +from modalities.models.nemotron.nemotron_mlp import SquaredReLUMLP + + +class _ResidualLayer(nn.Module): + """Base class implementing the shared pre-norm residual structure.""" + + def __init__(self, norm: nn.Module): + """ + Initializes the residual layer. + + Args: + norm (nn.Module): The pre-normalization module, owned exclusively by this layer. + """ + super().__init__() + self.norm = norm + + def _operator(self, x: torch.Tensor) -> torch.Tensor: + """ + Applies the layer's operator to the normalized input. + + Args: + x (torch.Tensor): The normalized input of shape ``(B, L, n_embd)``. + + Returns: + torch.Tensor: The operator output of shape ``(B, L, n_embd)``. + """ + raise NotImplementedError + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass: a pre-norm residual application of the layer's operator. + + Args: + x (torch.Tensor): Input of shape ``(B, L, n_embd)``. + + Returns: + torch.Tensor: Output of shape ``(B, L, n_embd)``. + """ + return x + self._operator(self.norm(x)) + + +class Mamba2Layer(_ResidualLayer): + """A residual layer whose operator is a Mamba-2 mixer.""" + + def __init__(self, norm: nn.Module, mixer: Mamba2Mixer): + """ + Initializes the Mamba2Layer. + + Args: + norm (nn.Module): The pre-normalization module. + mixer (Mamba2Mixer): The Mamba-2 mixer. + """ + super().__init__(norm=norm) + self.mixer = mixer + + def _operator(self, x: torch.Tensor) -> torch.Tensor: + return self.mixer(x) + + +class NemotronAttentionLayer(_ResidualLayer): + """A residual layer whose operator is grouped-query causal self-attention.""" + + def __init__(self, norm: nn.Module, attn: NemotronSelfAttention): + """ + Initializes the NemotronAttentionLayer. + + Args: + norm (nn.Module): The pre-normalization module. + attn (NemotronSelfAttention): The self-attention module. + """ + super().__init__(norm=norm) + self.attn = attn + + def _operator(self, x: torch.Tensor) -> torch.Tensor: + return self.attn(x) + + +class NemotronMoELayer(_ResidualLayer): + """A residual layer whose operator is a sparse mixture-of-experts feed-forward block.""" + + def __init__(self, norm: nn.Module, moe: MoE): + """ + Initializes the NemotronMoELayer. + + Args: + norm (nn.Module): The pre-normalization module. + moe (MoE): The mixture-of-experts block. + """ + super().__init__(norm=norm) + self.moe = moe + + def _operator(self, x: torch.Tensor) -> torch.Tensor: + return self.moe(x) + + +class NemotronMLPLayer(_ResidualLayer): + """A residual layer whose operator is a dense squared-ReLU feed-forward network.""" + + def __init__(self, norm: nn.Module, mlp: SquaredReLUMLP): + """ + Initializes the NemotronMLPLayer. + + Args: + norm (nn.Module): The pre-normalization module. + mlp (SquaredReLUMLP): The dense feed-forward network. + """ + super().__init__(norm=norm) + self.mlp = mlp + + def _operator(self, x: torch.Tensor) -> torch.Tensor: + return self.mlp(x) diff --git a/src/modalities/models/nemotron/nemotron_mlp.py b/src/modalities/models/nemotron/nemotron_mlp.py new file mode 100644 index 000000000..688368ec8 --- /dev/null +++ b/src/modalities/models/nemotron/nemotron_mlp.py @@ -0,0 +1,44 @@ +"""Non-gated squared-ReLU feed-forward network. + +Nemotron models use ``ReLU(x)^2`` rather than a gated activation such as SwiGLU. This halves the +number of matrices per feed-forward block (no gate projection), which is what makes a granular +128-expert MoE affordable. +""" + +import torch +import torch.nn as nn + +from modalities.models.components.moe.experts import squared_relu + + +class SquaredReLUMLP(nn.Module): + """ + A two-layer feed-forward network with a squared ReLU activation. + + Used both for dense feed-forward layers and as the body of the shared expert in an MoE layer. + """ + + def __init__(self, n_embd: int, ffn_hidden: int, bias: bool = False): + """ + Initializes the SquaredReLUMLP. + + Args: + n_embd (int): The model dimension. + ffn_hidden (int): The hidden dimension. + bias (bool): Whether the projections use a bias. Nemotron uses no biases anywhere. + """ + super().__init__() + self.c_fc = nn.Linear(in_features=n_embd, out_features=ffn_hidden, bias=bias) + self.c_proj = nn.Linear(in_features=ffn_hidden, out_features=n_embd, bias=bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass of the SquaredReLUMLP. + + Args: + x (torch.Tensor): Input of shape ``(..., n_embd)``. + + Returns: + torch.Tensor: Output of shape ``(..., n_embd)``. + """ + return self.c_proj(squared_relu(self.c_fc(x))) diff --git a/src/modalities/models/nemotron/nemotron_model.py b/src/modalities/models/nemotron/nemotron_model.py new file mode 100644 index 000000000..bb796d8d3 --- /dev/null +++ b/src/modalities/models/nemotron/nemotron_model.py @@ -0,0 +1,326 @@ +"""Nemotron-style hybrid Mamba-Transformer language model. + +The architecture (Nemotron-3 Nano 30B-A3B, arXiv:2512.20848) interleaves Mamba-2 mixers, +grouped-query attention and sparse mixture-of-experts feed-forward layers according to a layer +pattern string. Attention appears in only a handful of layers; the Mamba-2 layers carry both the +bulk of the sequence mixing and all positional information, which is why the model uses no +positional embeddings at all. + +The module tree deliberately mirrors GPT2's (``transformer.wte`` / ``transformer.h`` / +``transformer.lm_head_norm`` / ``transformer.lm_head``) so that the existing Modalities components +for activation checkpointing, FSDP wrapping, pipeline splitting and the chunked loss apply +unchanged. ``transformer.h`` is an ``nn.ModuleDict`` because the activation checkpointing component +requires one. +""" + +import logging +from typing import Annotated, Optional, overload + +import torch +import torch.nn as nn +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from modalities.config.pydantic_if_types import PydanticNemotronLayerSpecIFType +from modalities.models.components.norms import NormWrapperConfig +from modalities.models.model import NNModel +from modalities.models.nemotron.layer_pattern import LayerSymbol, parse_layer_pattern +from modalities.models.nemotron.nemotron_layer_specs import NemotronLayerSpecIF + +logger = logging.getLogger(__name__) + +# Matrix dimensions should be multiples of this for efficient tensor-core utilization, see +# https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/ +_TENSOR_CORE_ALIGNMENT = 128 + + +class NemotronLLMConfig(BaseModel): + """ + Configuration of the :class:`NemotronLLM` model. + + Attributes: + sample_key (str): Key under which the input token ids are found. + prediction_key (str): Key under which the logits are stored. + aux_loss_key (str | None): Key under which the summed MoE auxiliary loss is stored. When + None, the auxiliary loss is not exposed in the model output. + sequence_length (int): Maximum supported sequence length. + vocab_size (int): Vocabulary size. + n_embd (int): Model dimension. + n_layer (int): Number of layers. Must equal the length of ``layer_pattern``. + layer_pattern (str): One symbol per layer, e.g. ``"MEM*E"``. + layer_specs (dict[str, NemotronLayerSpecIF]): Maps a layer pattern symbol to the builder + that produces layers of that type. Must cover every symbol used in the pattern. + lm_head_norm_config (NormWrapperConfig): Normalization applied before the language model + head. + use_weight_tying (bool): Whether to tie the input embedding and the output projection. + Nemotron unties them. + use_meta_device (bool): Whether to build the model on the meta device. + enforce_tensor_core_alignment (bool): Whether to require that ``n_embd`` and ``vocab_size`` + are multiples of 128. + """ + + sample_key: str + prediction_key: str + aux_loss_key: Optional[str] = None + sequence_length: Annotated[int, Field(strict=True, ge=1)] + vocab_size: Annotated[int, Field(strict=True, ge=1)] + n_embd: Annotated[int, Field(strict=True, ge=1)] + n_layer: Annotated[int, Field(strict=True, ge=1)] + layer_pattern: str + layer_specs: dict[str, PydanticNemotronLayerSpecIFType] + lm_head_norm_config: NormWrapperConfig + use_weight_tying: bool = False + use_meta_device: Optional[bool] = False + enforce_tensor_core_alignment: bool = True + + # Avoid the pydantic warning about the protected 'model_' namespace. + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="after") + def _validate(self) -> "NemotronLLMConfig": + layer_symbols = parse_layer_pattern(self.layer_pattern) + if len(layer_symbols) != self.n_layer: + raise ValueError( + f"n_layer ({self.n_layer}) does not match the length of layer_pattern " + f"('{self.layer_pattern}' has {len(layer_symbols)} layers)." + ) + + # Every symbol appearing in the pattern needs a spec, and every provided spec must declare + # the symbol it was registered under. Both mistakes are easy to make in YAML. + for symbol, spec in self.layer_specs.items(): + parsed_symbol = LayerSymbol(symbol) + if spec.symbol != parsed_symbol: + raise ValueError( + f"layer_specs['{symbol}'] is a spec for layer type '{spec.symbol.value}', " + f"not '{parsed_symbol.value}'." + ) + missing = sorted({symbol.value for symbol in layer_symbols} - set(self.layer_specs)) + if missing: + raise ValueError( + f"layer_pattern uses the layer types {missing} but layer_specs only provides " + f"{sorted(self.layer_specs)}." + ) + unused = sorted(set(self.layer_specs) - {symbol.value for symbol in layer_symbols}) + if unused: + logger.warning("layer_specs provides unused layer types %s for pattern '%s'.", unused, self.layer_pattern) + + if self.enforce_tensor_core_alignment: + for name, value in (("n_embd", self.n_embd), ("vocab_size", self.vocab_size)): + if value % _TENSOR_CORE_ALIGNMENT != 0: + raise ValueError( + f"{name} with value {value} should be divisible by {_TENSOR_CORE_ALIGNMENT} for " + f"efficient training. Set enforce_tensor_core_alignment=False to override." + ) + return self + + +class NemotronLLM(NNModel): + """A hybrid Mamba-Transformer decoder-only language model.""" + + def __init__( + self, + sample_key: str, + prediction_key: str, + sequence_length: int, + vocab_size: int, + n_embd: int, + n_layer: int, + layer_pattern: str, + layer_specs: dict[str, NemotronLayerSpecIF], + lm_head_norm_config: NormWrapperConfig, + use_weight_tying: bool = False, + aux_loss_key: Optional[str] = None, + ): + """ + Initializes the NemotronLLM. + + Args: + sample_key (str): Key under which the input token ids are found. + prediction_key (str): Key under which the logits are stored. + sequence_length (int): Maximum supported sequence length. + vocab_size (int): Vocabulary size. + n_embd (int): Model dimension. + n_layer (int): Number of layers. + layer_pattern (str): One symbol per layer. + layer_specs (dict[str, NemotronLayerSpecIF]): Builders keyed by layer pattern symbol. + lm_head_norm_config (NormWrapperConfig): Normalization before the language model head. + use_weight_tying (bool): Whether to tie the embedding and the output projection. + aux_loss_key (str | None): Key under which the summed MoE auxiliary loss is exposed. + """ + weight_decay_groups = { + "linear": [ + r"\.attn\.", + r"\.mixer\.in_proj", + r"\.mixer\.out_proj", + r"\.mlp\.", + r"\.shared_experts\.", + r"\.lm_head\.weight", + ], + "experts": [r"\.experts\.w1", r"\.experts\.w2"], + "router": [r"\.router\.gate"], + "embedding": [r"\.wte\."], + "layernorm": [r"\.norm\.", r"\.lm_head_norm\."], + "ssm": [r"\.A_log", r"\.D$", r"\.dt_bias", r"\.conv1d_weight", r"\.conv1d_bias"], + } + super().__init__(weight_decay_groups=weight_decay_groups) + + self.sample_key = sample_key + self.prediction_key = prediction_key + self.aux_loss_key = aux_loss_key + self.sequence_length = sequence_length + self.vocab_size = vocab_size + self.n_embd = n_embd + self.n_layer = n_layer + self.layer_pattern = layer_pattern + # When True, forward returns the post-norm hidden states instead of logits so that a + # memory-efficient loss (e.g. ChunkedCLMCrossEntropyLoss) can apply the head in chunks. + self._skip_lm_head = False + + layer_symbols = parse_layer_pattern(layer_pattern) + if len(layer_symbols) != n_layer: + raise ValueError( + f"n_layer ({n_layer}) does not match the length of layer_pattern " + f"('{layer_pattern}' has {len(layer_symbols)} layers)." + ) + specs_by_symbol = {LayerSymbol(symbol): spec for symbol, spec in layer_specs.items()} + missing = sorted({symbol.value for symbol in layer_symbols} - {s.value for s in specs_by_symbol}) + if missing: + raise ValueError(f"No layer spec provided for layer types {missing}.") + + self.transformer = nn.ModuleDict( + dict( + wte=nn.Embedding(num_embeddings=vocab_size, embedding_dim=n_embd), + # A ModuleDict (rather than a ModuleList) is required by the Modalities activation + # checkpointing component and matches the GPT2 layout. + h=nn.ModuleDict( + { + str(layer_idx): specs_by_symbol[symbol].build(layer_idx=layer_idx) + for layer_idx, symbol in enumerate(layer_symbols) + } + ), + lm_head_norm=lm_head_norm_config.build(), + lm_head=nn.Linear(in_features=n_embd, out_features=vocab_size, bias=False), + ) + ) + + if use_weight_tying: + self.transformer.wte.weight = self.transformer.lm_head.weight + + @property + def lm_head(self) -> nn.Module: + """The language-model head. Exposed so a memory-efficient loss can apply it chunk-by-chunk.""" + return self.transformer.lm_head + + def set_skip_lm_head(self, skip: bool) -> None: + """ + Toggles whether forward returns post-norm hidden states instead of logits. + + Args: + skip (bool): True to return hidden states, False to return logits. + """ + self._skip_lm_head = skip + + @property + def has_tied_word_embeddings(self) -> bool: + """ + Whether the token embedding and the output projection currently share a weight tensor. + + Under pipeline parallelism a stage may contain neither submodule, in which case there is no + tying to report. + + Returns: + bool: True if the weights are tied. + """ + if "wte" not in self.transformer or "lm_head" not in self.transformer: + return False + return self.transformer.wte.weight is self.transformer.lm_head.weight + + def get_moe_layers(self) -> list[nn.Module]: + """ + Returns every mixture-of-experts block contained in this model (or pipeline stage). + + Returns: + list[nn.Module]: The MoE blocks, in model order. + """ + # Imported here to keep the module import graph acyclic at definition time. + from modalities.models.components.moe.moe import MoE + + return [module for module in self.modules() if isinstance(module, MoE)] + + def get_aux_loss(self) -> Optional[torch.Tensor]: + """ + Sums the auxiliary load-balancing losses recorded by the MoE layers in the last forward pass. + + Returns: + torch.Tensor | None: The summed auxiliary loss, or None if no MoE layer produced one + (either because no MoE layers exist on this stage or because the coefficient is 0). + """ + losses = [moe.last_aux_loss for moe in self.get_moe_layers() if moe.last_aux_loss is not None] + if not losses: + return None + return torch.stack(losses).sum() + + @overload + def forward(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + ... + + @overload + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + ... + + def forward(self, inputs: dict[str, torch.Tensor] | torch.Tensor) -> dict[str, torch.Tensor] | torch.Tensor: + """ + Forward pass of the model. + + Args: + inputs (dict[str, torch.Tensor] | torch.Tensor): Either a dict containing the input + token ids under ``sample_key``, or the token id tensor directly (used by pipeline + parallelism, which passes raw tensors between stages). + + Returns: + dict[str, torch.Tensor] | torch.Tensor: A dict with the logits under ``prediction_key`` + (and the auxiliary loss under ``aux_loss_key`` if configured), or the logits tensor. + """ + if not isinstance(inputs, dict): + return self.forward_impl(inputs) + + logits = self.forward_impl(inputs[self.sample_key]) + # Insertion order matters: InferenceResultBatch derives its length from the first entry, + # so the token-shaped logits have to come first. + predictions = {self.prediction_key: logits} + if self.aux_loss_key is not None: + aux_loss = self.get_aux_loss() + if aux_loss is not None: + predictions[self.aux_loss_key] = aux_loss + return predictions + + def forward_impl(self, inputs: torch.Tensor) -> torch.Tensor: + """ + Forward pass implementation operating on plain tensors. + + Every submodule access is guarded by ``hasattr`` so that a pipeline stage holding only a + subset of the model still works. + + Args: + inputs (torch.Tensor): Input token ids of shape ``(B, L)``, or hidden states of shape + ``(B, L, n_embd)`` on a non-first pipeline stage. + + Returns: + torch.Tensor: Logits of shape ``(B, L, vocab_size)``, or hidden states if the head is + skipped or absent. + """ + seq_len = inputs.size(1) + if seq_len > self.sequence_length: + raise ValueError( + f"Cannot forward sequence of length {seq_len}; the model's maximum input sequence " + f"length is {self.sequence_length}." + ) + + h = self.transformer.wte(inputs) if hasattr(self.transformer, "wte") else inputs + + for layer_idx in self.transformer.h: + h = self.transformer.h[layer_idx](h) + + h = self.transformer.lm_head_norm(h) if hasattr(self.transformer, "lm_head_norm") else h + if self._skip_lm_head: + return h + return self.transformer.lm_head(h) if hasattr(self.transformer, "lm_head") else h diff --git a/src/modalities/models/nemotron/nemotron_model_factory.py b/src/modalities/models/nemotron/nemotron_model_factory.py new file mode 100644 index 000000000..909bb4dba --- /dev/null +++ b/src/modalities/models/nemotron/nemotron_model_factory.py @@ -0,0 +1,83 @@ +"""Factory for the Nemotron hybrid Mamba-Transformer model.""" + +from typing import Optional + +import torch + +from modalities.models.components.norms import NormWrapperConfig +from modalities.models.nemotron.nemotron_layer_specs import NemotronLayerSpecIF +from modalities.models.nemotron.nemotron_model import NemotronLLM + + +class NemotronModelFactory: + """Factory class creating Nemotron models.""" + + @staticmethod + def get_nemotron_model( + sample_key: str, + prediction_key: str, + sequence_length: int, + vocab_size: int, + n_embd: int, + n_layer: int, + layer_pattern: str, + layer_specs: dict[str, NemotronLayerSpecIF], + lm_head_norm_config: NormWrapperConfig | dict, + use_weight_tying: bool = False, + aux_loss_key: Optional[str] = None, + use_meta_device: Optional[bool] = False, + enforce_tensor_core_alignment: bool = True, + ) -> NemotronLLM: + """ + Creates a NemotronLLM, optionally on the meta device. + + Args: + sample_key (str): Key under which the input token ids are found. + prediction_key (str): Key under which the logits are stored. + sequence_length (int): Maximum supported sequence length. + vocab_size (int): Vocabulary size. + n_embd (int): Model dimension. + n_layer (int): Number of layers. + layer_pattern (str): One symbol per layer. + layer_specs (dict[str, NemotronLayerSpecIF]): Builders keyed by layer pattern symbol. + lm_head_norm_config (NormWrapperConfig | dict): Normalization before the language model + head. A plain dict is accepted so that the factory can also be called directly, + outside the component factory that would normally validate it. + use_weight_tying (bool): Whether to tie the embedding and the output projection. + aux_loss_key (str | None): Key under which the summed MoE auxiliary loss is exposed. + use_meta_device (bool): Whether to build the model on the meta device. Materialization + and initialization then happen in ``ModelFactory.get_weight_initialized_model``. + enforce_tensor_core_alignment (bool): Validated in the config; accepted here so that the + config and the factory signature stay in sync. + + Raises: + ValueError: If weight tying is combined with meta device initialization. + + Returns: + NemotronLLM: The constructed model. + """ + del enforce_tensor_core_alignment # validated by NemotronLLMConfig + if not isinstance(lm_head_norm_config, NormWrapperConfig): + lm_head_norm_config = NormWrapperConfig.model_validate(lm_head_norm_config) + config = dict( + sample_key=sample_key, + prediction_key=prediction_key, + sequence_length=sequence_length, + vocab_size=vocab_size, + n_embd=n_embd, + n_layer=n_layer, + layer_pattern=layer_pattern, + layer_specs=layer_specs, + lm_head_norm_config=lm_head_norm_config, + use_weight_tying=use_weight_tying, + aux_loss_key=aux_loss_key, + ) + if use_meta_device and use_weight_tying: + raise ValueError( + "Weight tying is not supported on the meta device. Set either use_meta_device=False " + "or use_weight_tying=False. See https://github.com/Modalities/modalities/issues/357" + ) + if use_meta_device: + with torch.device("meta"): + return NemotronLLM(**config) + return NemotronLLM(**config) diff --git a/src/modalities/models/nemotron/nemotron_stages_generator.py b/src/modalities/models/nemotron/nemotron_stages_generator.py new file mode 100644 index 000000000..68cf1d2a9 --- /dev/null +++ b/src/modalities/models/nemotron/nemotron_stages_generator.py @@ -0,0 +1,246 @@ +"""Pipeline stage generation for the Nemotron hybrid model. + +The GPT2 stages generator assumes every layer costs the same. That is a poor assumption for a hybrid +MoE model: a mixture-of-experts layer performs far more work per token than a Mamba-2 layer, and an +attention layer sits somewhere in between. Splitting such a stack into equal *counts* of layers +produces badly unbalanced pipeline stages, and the slowest stage sets the throughput of the whole +pipeline. This generator therefore assigns a per-layer-type computational weight. +""" + +import math +from typing import Annotated + +from pydantic import BaseModel, Field, model_validator + +from modalities.models.nemotron.layer_pattern import LayerSymbol, parse_layer_pattern +from modalities.models.parallelism.pipeline_parallelism import StagesGenerator + +# Relative computational weights per layer type. The pipeline packer works with integers, so the +# weights are expressed on a scale where the cheapest layer type is 1. +DEFAULT_LAYER_WEIGHTS: dict[LayerSymbol, int] = { + LayerSymbol.MAMBA: 2, + LayerSymbol.ATTENTION: 1, + LayerSymbol.MOE: 3, + LayerSymbol.MLP: 1, +} + + +class NemotronStagesGeneratorConfig(BaseModel): + """ + Configuration of :class:`NemotronStagesGenerator`. + + Attributes: + layer_pattern (str): The model's layer pattern, one symbol per layer. + input_layer_equivalence (int): Computational weight of the embedding stage, expressed in + units of the per-layer weights. + output_layer_equivalence (int): Computational weight of the head stage. + layer_weights (dict[str, int] | None): Optional override of the per-layer-type weights. + """ + + layer_pattern: str + input_layer_equivalence: Annotated[int, Field(strict=True, ge=1)] = 2 + output_layer_equivalence: Annotated[int, Field(strict=True, ge=1)] = 2 + layer_weights: dict[str, Annotated[int, Field(strict=True, ge=1)]] | None = None + + @model_validator(mode="after") + def _validate(self) -> "NemotronStagesGeneratorConfig": + layer_symbols = parse_layer_pattern(self.layer_pattern) + if self.layer_weights is not None: + missing = sorted({symbol.value for symbol in layer_symbols} - set(self.layer_weights)) + if missing: + raise ValueError(f"layer_weights is missing an entry for the layer types {missing}.") + return self + + +class NemotronStagesGenerator(StagesGenerator): + """Generates fully qualified module names per pipeline stage, weighted by layer type.""" + + def __init__( + self, + layer_pattern: str, + input_layer_equivalence: int = 2, + output_layer_equivalence: int = 2, + layer_weights: dict[str, int] | None = None, + ): + """ + Initializes the NemotronStagesGenerator. + + Args: + layer_pattern (str): The model's layer pattern. + input_layer_equivalence (int): Computational weight of the embedding stage. + output_layer_equivalence (int): Computational weight of the head stage. + layer_weights (dict[str, int] | None): Optional override of the per-layer-type weights. + """ + self._layer_symbols = parse_layer_pattern(layer_pattern) + super().__init__( + num_model_layers=len(self._layer_symbols), + input_layer_equivalence=input_layer_equivalence, + output_layer_equivalence=output_layer_equivalence, + ) + if layer_weights is None: + self._layer_weights = dict(DEFAULT_LAYER_WEIGHTS) + else: + self._layer_weights = {LayerSymbol(symbol): weight for symbol, weight in layer_weights.items()} + + def _get_potential_split_points(self) -> list[tuple[list[str], int]]: + """ + Returns the candidate pipeline split points with their computational weights. + + Returns: + list[tuple[list[str], int]]: Fully qualified module names per split point, paired with + the relative cost of that split point. + """ + return [ + (["transformer.wte"], self._input_layer_equivalence), + *[ + ([f"transformer.h.{layer_idx}"], self._layer_weights[symbol]) + for layer_idx, symbol in enumerate(self._layer_symbols) + ], + (["transformer.lm_head_norm", "transformer.lm_head"], self._output_layer_equivalence), + ] + + def get_stages(self, num_layers_per_stage: int, pp_dims: int) -> list[list[str]]: + """ + Partitions the model into pipeline stages, balancing computational cost. + + The base-class implementation packs split points greedily against a fixed per-stage weight + cap. With uniform layer weights (as in GPT2) that always consumes every split point, but + with the per-layer-type weights used here a greedy pass can exhaust its stage budget before + reaching the end and silently drop the trailing modules - producing a model with no output + layer on any stage. This override therefore uses a partitioner that is guaranteed to assign + every module to exactly one stage. + + The algorithm finds the smallest per-stage weight cap for which a left-to-right greedy pass + fits within the requested number of stages (binary search over the cap), and then splits the + heaviest stages further until the requested count is reached. That yields contiguous, + complete, and near-optimally balanced stages. + + Args: + num_layers_per_stage (int): Target number of layers per stage. Together with the input + and output equivalences this determines the number of virtual stages. + pp_dims (int): Pipeline parallel degree. The number of virtual stages must be a multiple + of this. + + Raises: + ValueError: If the number of virtual stages is not a multiple of ``pp_dims``, or if there + are fewer split points than requested stages. + + Returns: + list[list[str]]: Fully qualified module names per stage, in model order. + """ + if num_layers_per_stage < 1: + raise ValueError(f"num_layers_per_stage must be at least 1, got {num_layers_per_stage}.") + + num_virtual_stages = math.ceil( + (self._num_model_layers + self._input_layer_equivalence + self._output_layer_equivalence) + / num_layers_per_stage + ) + if num_virtual_stages % pp_dims != 0: + raise ValueError( + f"Number of virtual stages {num_virtual_stages} is not divisible by parallel dimensions " + f"{pp_dims}. For reference: {self._num_model_layers=} {self._input_layer_equivalence=} " + f"{self._output_layer_equivalence=} {num_layers_per_stage=}" + ) + + split_points = self._get_potential_split_points() + if num_virtual_stages > len(split_points): + raise ValueError( + f"Cannot build {num_virtual_stages} pipeline stages from only {len(split_points)} " + f"split points. Increase num_layers_per_stage or reduce the pipeline degree." + ) + + groups = _partition_contiguous(split_points, num_parts=num_virtual_stages) + return [[fqn for fqns, _ in group for fqn in fqns] for group in groups] + + +def _greedy_pack(split_points: list[tuple[list[str], int]], weight_cap: int) -> list[list[tuple[list[str], int]]]: + """ + Packs split points left to right into contiguous groups, each at most ``weight_cap`` heavy. + + Args: + split_points (list[tuple[list[str], int]]): The split points with their weights. + weight_cap (int): Maximum weight per group. Must be at least the heaviest single split point. + + Returns: + list[list[tuple[list[str], int]]]: The resulting groups. Every split point is assigned. + """ + groups: list[list[tuple[list[str], int]]] = [] + current: list[tuple[list[str], int]] = [] + current_weight = 0 + for split_point in split_points: + weight = split_point[1] + if current and current_weight + weight > weight_cap: + groups.append(current) + current, current_weight = [], 0 + current.append(split_point) + current_weight += weight + if current: + groups.append(current) + return groups + + +def _best_binary_split(group: list[tuple[list[str], int]]) -> int: + """ + Finds the index at which splitting a group minimizes the weight of its heavier half. + + Args: + group (list[tuple[list[str], int]]): The group to split. Must contain at least two entries. + + Returns: + int: The split index, in ``[1, len(group) - 1]``. + """ + weights = [weight for _, weight in group] + total = sum(weights) + best_index, best_cost = 1, None + prefix = 0 + for index in range(1, len(group)): + prefix += weights[index - 1] + cost = max(prefix, total - prefix) + if best_cost is None or cost < best_cost: + best_index, best_cost = index, cost + return best_index + + +def _partition_contiguous( + split_points: list[tuple[list[str], int]], num_parts: int +) -> list[list[tuple[list[str], int]]]: + """ + Partitions split points into exactly ``num_parts`` contiguous, non-empty, balanced groups. + + Args: + split_points (list[tuple[list[str], int]]): The split points with their weights, in order. + num_parts (int): The exact number of groups to produce. + + Raises: + ValueError: If there are fewer split points than requested groups. + + Returns: + list[list[tuple[list[str], int]]]: The groups, covering every split point exactly once. + """ + if num_parts > len(split_points): + raise ValueError(f"Cannot partition {len(split_points)} split points into {num_parts} groups.") + if num_parts == 1: + return [list(split_points)] + + weights = [weight for _, weight in split_points] + low, high = max(weights), sum(weights) + feasible_cap = high + while low <= high: + candidate = (low + high) // 2 + if len(_greedy_pack(split_points, weight_cap=candidate)) <= num_parts: + feasible_cap = candidate + high = candidate - 1 + else: + low = candidate + 1 + + groups = _greedy_pack(split_points, weight_cap=feasible_cap) + # The binary search only guarantees "at most num_parts" groups. Split the heaviest splittable + # group until the requested count is reached; pipeline parallelism needs exactly this many. + while len(groups) < num_parts: + splittable = [index for index, group in enumerate(groups) if len(group) > 1] + heaviest = max(splittable, key=lambda index: sum(weight for _, weight in groups[index])) + group = groups.pop(heaviest) + split_index = _best_binary_split(group) + groups.insert(heaviest, group[split_index:]) + groups.insert(heaviest, group[:split_index]) + return groups diff --git a/src/modalities/nn/model_initialization/initialization_routines.py b/src/modalities/nn/model_initialization/initialization_routines.py index 1f785f562..4482e2cc3 100644 --- a/src/modalities/nn/model_initialization/initialization_routines.py +++ b/src/modalities/nn/model_initialization/initialization_routines.py @@ -18,6 +18,34 @@ class MultiDeviceGeneratorPolicy(str, Enum): ERROR = "error" +# Wrappers that insert themselves into a parameter's fully qualified name without changing which +# logical parameter it is. The initialization filters are written against the plain model FQNs, so +# these prefixes are stripped before matching. Without this, applying activation checkpointing (or +# torch.compile) before initialization would silently prevent every per-layer regex from matching, +# leaving the model with its default rather than its configured initialization. +_FQN_WRAPPER_PREFIXES = ( + "_orig_mod.", # torch.compile + "_checkpoint_wrapped_module.", # activation checkpointing + "_fsdp_wrapped_module.", # FSDP1 +) + + +def normalize_parameter_name(parameter_name: str) -> str: + """ + Removes wrapper prefixes from a parameter's fully qualified name. + + Args: + parameter_name (str): The fully qualified parameter name, possibly containing wrapper + segments such as ``_checkpoint_wrapped_module.``. + + Returns: + str: The name as it would appear on the unwrapped model. + """ + for prefix in _FQN_WRAPPER_PREFIXES: + parameter_name = parameter_name.replace(prefix, "") + return parameter_name + + class PlainInitializationConfig(BaseModel): mean: float std: Annotated[float, Field(strict=True, ge=0.0)] | str # can be float or "auto" @@ -85,9 +113,9 @@ def initialize_in_place(self, model: nn.Module): weight_regexes = self.parameter_name_regexes.weights bias_regexes = self.parameter_name_regexes.biases or [] for parameter_name, p in model.named_parameters(): - parameter_name = parameter_name.replace( - "_orig_mod.", "" - ) # remove FQN modification from torch.compile if present + # Strip FQN modifications introduced by torch.compile, activation checkpointing and + # FSDP1 so that the filters can be written against the plain model. + parameter_name = normalize_parameter_name(parameter_name) for weight_regex in weight_regexes: if re.fullmatch(weight_regex, parameter_name): nn.init.normal_(p, mean=self.mean, std=self.std, generator=self._get_generator(p)) diff --git a/src/modalities/nn/model_initialization/parameter_name_filters.py b/src/modalities/nn/model_initialization/parameter_name_filters.py index ebc8a23ec..b69085ee3 100644 --- a/src/modalities/nn/model_initialization/parameter_name_filters.py +++ b/src/modalities/nn/model_initialization/parameter_name_filters.py @@ -13,6 +13,7 @@ class WeightInitTypes(Enum): class SupportWeightInitModels(Enum): GPT2 = "gpt2" COCA = "coca" + NEMOTRON = "nemotron" class RegexFilter(BaseModel): @@ -74,4 +75,52 @@ class RegexFilter(BaseModel): WeightInitTypes.SCALED: RegexFilter(weights=[], biases=[]), WeightInitTypes.SCALED_EMBED: RegexFilter(weights=[], biases=[]), }, + SupportWeightInitModels.NEMOTRON: { + # Only the linear and embedding weights of the hybrid Mamba-Transformer are drawn from a + # normal distribution. The state space parameters (A_log, D, dt_bias, conv1d_*) follow + # their own distributions, defined in Mamba2Mixer.reset_parameters, and the router bias + # must stay at zero; none of them may be matched here or they would be overwritten. + WeightInitTypes.PLAIN: RegexFilter( + weights=[ + # Mamba-2 mixer projections + r"transformer\.h\.\d+\.mixer\.(in_proj|out_proj)\.weight", + # attention projections + r"transformer\.h\.\d+\.attn\.(q_attn|k_attn|v_attn|c_proj)\.weight", + # dense feed-forward + r"transformer\.h\.\d+\.mlp\.(c_fc|c_proj)\.weight", + # mixture-of-experts: router gate, routed experts, shared experts + r"transformer\.h\.\d+\.moe\.router\.gate\.weight", + r"transformer\.h\.\d+\.moe\.experts\.w[12]", + r"transformer\.h\.\d+\.moe\.shared_experts\.(c_fc|c_proj)\.weight", + # embeddings and (untied) language model head + r"transformer\.wte\.weight", + r"transformer\.lm_head\.weight", + ], + biases=[ + # Nemotron uses no biases, but they remain configurable per component. + r"transformer\.h\.\d+\.mixer\.(in_proj|out_proj)\.bias", + r"transformer\.h\.\d+\.attn\.(q_attn|k_attn|v_attn|c_proj)\.bias", + r"transformer\.h\.\d+\.mlp\.(c_fc|c_proj)\.bias", + r"transformer\.h\.\d+\.moe\.shared_experts\.(c_fc|c_proj)\.bias", + ], + ), + # scaled: down-scale every projection that writes into the residual stream, so that the + # variance contributed by the residual branches does not grow with depth + # (https://arxiv.org/abs/2312.16903). + WeightInitTypes.SCALED: RegexFilter( + weights=[ + r"transformer\.h\.\d+\.mixer\.out_proj\.weight", + r"transformer\.h\.\d+\.attn\.c_proj\.weight", + r"transformer\.h\.\d+\.mlp\.c_proj\.weight", + r"transformer\.h\.\d+\.moe\.experts\.w2", + r"transformer\.h\.\d+\.moe\.shared_experts\.c_proj\.weight", + ] + ), + WeightInitTypes.SCALED_EMBED: RegexFilter( + weights=[ + r"transformer\.wte\.weight", + r"transformer\.lm_head\.weight", + ] + ), + }, } diff --git a/src/modalities/registry/components.py b/src/modalities/registry/components.py index d9d0a034c..6bc548d93 100644 --- a/src/modalities/registry/components.py +++ b/src/modalities/registry/components.py @@ -52,6 +52,7 @@ LinearWarmupCosineAnnealingLRSchedulerConfig, LLMDataLoaderConfig, MemMapDatasetConfig, + MoELoadBalancedOptimizerConfig, OneCycleLRSchedulerConfig, PackedMemMapDatasetContinuousConfig, PackedMemMapDatasetMegatronConfig, @@ -92,11 +93,31 @@ RMSLayerNorm, RMSLayerNormConfig, ) +from modalities.models.components.moe.load_balancing import MoEBalancing +from modalities.models.components.moe.moe_losses import ( + MoEAuxLoss, + MoEAuxLossConfig, + WeightedSumLoss, + WeightedSumLossConfig, +) from modalities.models.gpt2.collator import GPT2LLMCollateFn from modalities.models.gpt2.gpt2_model import GPT2LLMConfig from modalities.models.gpt2.llama3_like_initialization import Llama3Initializer, Llama3InitializerConfig from modalities.models.huggingface.huggingface_model import HuggingFacePretrainedModel, HuggingFacePretrainedModelConfig from modalities.models.model_factory import GPT2ModelFactory, ModelFactory +from modalities.models.nemotron.nemotron_layer_specs import ( + Mamba2LayerSpec, + Mamba2LayerSpecConfig, + NemotronAttentionLayerSpec, + NemotronAttentionLayerSpecConfig, + NemotronMLPLayerSpec, + NemotronMLPLayerSpecConfig, + NemotronMoELayerSpec, + NemotronMoELayerSpecConfig, +) +from modalities.models.nemotron.nemotron_model import NemotronLLMConfig +from modalities.models.nemotron.nemotron_model_factory import NemotronModelFactory +from modalities.models.nemotron.nemotron_stages_generator import NemotronStagesGenerator, NemotronStagesGeneratorConfig from modalities.models.parallelism.pipeline_parallelism import ComponentSelectorFromPipeline, PipelineFactory from modalities.models.parallelism.pipeline_parallelism_configs import ( ComponentSelectorFromPipelineConfig, @@ -132,6 +153,7 @@ from modalities.utils.debugging_configs import DebuggingConfig, NaNHookConfig, PrintForwardHookConfig from modalities.utils.maybe_list_parameter import MaybeListDecorator, maybe_list_parameter from modalities.utils.mfu import GPT2MFUCalculator +from modalities.utils.nemotron_mfu import NemotronMFUCalculator, NemotronMFUCalculatorConfig from modalities.utils.number_conversion import ( LocalNumBatchesFromNumSamplesConfig, LocalNumBatchesFromNumTokensConfig, @@ -221,6 +243,14 @@ class ComponentEntity: ), ComponentEntity("model", "compiled", maybe_model_list(ModelFactory.get_compiled_model), CompiledModelConfig), ComponentEntity("model", "coca", CoCa, CoCaConfig), + # Nemotron hybrid Mamba-Transformer + ComponentEntity("model", "nemotron", NemotronModelFactory.get_nemotron_model, NemotronLLMConfig), + # Nemotron layer specs. These are builders, not modules: the model calls build() once per + # layer position so that every layer gets its own parameters (see nemotron_layer_specs.py). + ComponentEntity("nemotron_layer_spec", "mamba2", Mamba2LayerSpec, Mamba2LayerSpecConfig), + ComponentEntity("nemotron_layer_spec", "attention", NemotronAttentionLayerSpec, NemotronAttentionLayerSpecConfig), + ComponentEntity("nemotron_layer_spec", "moe", NemotronMoELayerSpec, NemotronMoELayerSpecConfig), + ComponentEntity("nemotron_layer_spec", "mlp", NemotronMLPLayerSpec, NemotronMLPLayerSpecConfig), ComponentEntity( "model", "debugging_enriched", @@ -233,6 +263,9 @@ class ComponentEntity: ComponentEntity("pipeline", "builder", PipelineFactory.get_pipeline, PipelineConfig), # Pipeline Stages Generators ComponentEntity("stages_generator", "gpt2_stages_generator", GPT2LLMStagesGenerator, GPT2LLMStagesGeneratorConfig), + ComponentEntity( + "stages_generator", "nemotron_stages_generator", NemotronStagesGenerator, NemotronStagesGeneratorConfig + ), # Device mesh ComponentEntity("device_mesh", "default", get_device_mesh, DeviceMeshConfig), ComponentEntity("number_conversion", "parallel_degree", get_parallel_degree, ParallelDegreeConfig), @@ -251,6 +284,8 @@ class ComponentEntity: ), # losses ComponentEntity("loss", "clm_cross_entropy_loss", CLMCrossEntropyLoss, CLMCrossEntropyLossConfig), + ComponentEntity("loss", "moe_aux_loss", MoEAuxLoss, MoEAuxLossConfig), + ComponentEntity("loss", "weighted_sum", WeightedSumLoss, WeightedSumLossConfig), # optimizers ComponentEntity( "optimizer", "adam", maybe_model_list_for_optimizer(OptimizerFactory.get_adam), AdamOptimizerConfig @@ -264,6 +299,13 @@ class ComponentEntity: OptimizerFactory.get_fsdp1_checkpointed_optimizer_, FSDP1CheckpointedOptimizerConfig, ), + # Optimizer decorator: adds the auxiliary-loss-free MoE load balancing step to any optimizer. + ComponentEntity( + "optimizer", + "moe_load_balanced", + MoEBalancing.register_expert_bias_update_hook, + MoELoadBalancedOptimizerConfig, + ), # App state ComponentEntity("app_state", "raw", AppStateFactory.get_raw_app_state, RawAppStateConfig), ComponentEntity("app_state", "dcp", AppStateFactory.get_dcp_checkpointed_app_state_, DCPAppStateConfig), @@ -414,6 +456,7 @@ class ComponentEntity: ), # MFU calculators ComponentEntity("mfu_calculator", "gpt2", GPT2MFUCalculator, GPT2MFUCalculatorConfig), + ComponentEntity("mfu_calculator", "nemotron", NemotronMFUCalculator, NemotronMFUCalculatorConfig), # Number conversion ComponentEntity( "number_conversion", diff --git a/src/modalities/utils/nemotron_mfu.py b/src/modalities/utils/nemotron_mfu.py new file mode 100644 index 000000000..8d8bf9bfc --- /dev/null +++ b/src/modalities/utils/nemotron_mfu.py @@ -0,0 +1,173 @@ +"""Model FLOPs Utilization for the Nemotron hybrid Mamba-Transformer. + +The GPT2 estimate ``6 * num_params + 12 * n_layer * seq_len * n_embd`` is wrong for this +architecture on two counts: + +1. It counts *all* parameters. In a 128-expert MoE only the 6 routed experts a token actually visits + contribute FLOPs, so using total parameters would overstate the work by roughly an order of + magnitude and report an absurdly high MFU. +2. Its second term is the attention score/value cost, which assumes *every* layer attends. In + Nemotron-3 Nano only 6 of 52 layers do; the other 46 are Mamba-2 or MoE layers whose cost is + linear in the sequence length. + +This calculator therefore takes the number of *active* parameters and adds the quadratic attention +term only for the attention layers. +""" + +from typing import Annotated, Optional + +import torch +from pydantic import BaseModel, ConfigDict, Field + +from modalities.config.pydantic_if_types import PydanticPytorchModuleOrListType +from modalities.models.nemotron.layer_pattern import LayerSymbol, count_layers_by_type +from modalities.utils.mfu import MFUCalculatorABC + + +class NemotronMFUCalculatorConfig(BaseModel): + """ + Configuration of :class:`NemotronMFUCalculator`. + + Attributes: + layer_pattern (str): The model's layer pattern, used to count attention layers. + sequence_length (int): Training sequence length. + n_embd (int): Model dimension. + n_head_q (int): Number of query heads of the attention layers. + head_dim (int): Head dimension of the attention layers. + num_active_params (int): Number of parameters visited per token. See + :meth:`NemotronMFUCalculator.count_active_parameters`. + world_size (int): Number of ranks. + model_parts (nn.Module | list[nn.Module]): The wrapped model (or pipeline stages). + device_mesh (DeviceMesh | None): The device mesh, if any. + """ + + layer_pattern: str + sequence_length: Annotated[int, Field(strict=True, ge=1)] + n_embd: Annotated[int, Field(strict=True, ge=1)] + n_head_q: Annotated[int, Field(strict=True, ge=1)] + head_dim: Annotated[int, Field(strict=True, ge=1)] + num_active_params: Annotated[int, Field(strict=True, ge=1)] + world_size: Annotated[int, Field(strict=True, ge=1)] + model_parts: PydanticPytorchModuleOrListType + device_mesh: Optional[object] = None + + # avoid the pydantic warning about the protected 'model_' namespace + model_config = ConfigDict(protected_namespaces=(), arbitrary_types_allowed=True) + + +class NemotronMFUCalculator(MFUCalculatorABC): + """Computes the Model FLOPs Utilization of a hybrid Mamba-Transformer MoE model.""" + + def __init__( + self, + layer_pattern: str, + sequence_length: int, + n_embd: int, + n_head_q: int, + head_dim: int, + num_active_params: int, + world_size: int, + model_parts: torch.nn.Module | list[torch.nn.Module], + device_mesh: Optional[object] = None, + ): + """ + Initializes the NemotronMFUCalculator. + + Args: + layer_pattern (str): The model's layer pattern. + sequence_length (int): Training sequence length. + n_embd (int): Model dimension. + n_head_q (int): Number of query heads of the attention layers. + head_dim (int): Head dimension of the attention layers. + num_active_params (int): Number of parameters visited per token. + world_size (int): Number of ranks. + model_parts (nn.Module | list[nn.Module]): The wrapped model or pipeline stages. + device_mesh (DeviceMesh | None): The device mesh, if any. + """ + del device_mesh # only needed by the parameter-counting calculators + self._sequence_length = sequence_length + self._num_attention_layers = count_layers_by_type(layer_pattern)[LayerSymbol.ATTENTION] + self._theoretical_flops_per_token = NemotronMFUCalculator._get_theoretical_flops_per_token( + num_active_params=num_active_params, + num_attention_layers=self._num_attention_layers, + sequence_length=sequence_length, + n_head_q=n_head_q, + head_dim=head_dim, + ) + self._theoretical_gpu_peak_performance = MFUCalculatorABC._get_theoretical_gpu_peak_performance( + model_parts, world_size + ) + del n_embd # part of the public config for symmetry with the GPT2 calculator + + @staticmethod + def _get_theoretical_flops_per_token( + num_active_params: int, + num_attention_layers: int, + sequence_length: int, + n_head_q: int, + head_dim: int, + ) -> int: + """ + Estimates the forward-plus-backward FLOPs per token. + + The first term is the usual ``6 * active_params`` (2 for the forward matmuls, 4 for the + backward). The second term is the attention score and value matmuls, which are quadratic in + the sequence length and only incurred by the attention layers: + ``12 * num_attention_layers * seq_len * n_head_q * head_dim``. + + Args: + num_active_params (int): Parameters visited per token. + num_attention_layers (int): Number of attention layers. + sequence_length (int): Sequence length. + n_head_q (int): Number of query heads. + head_dim (int): Attention head dimension. + + Returns: + int: Estimated FLOPs per token. + """ + dense_flops = 6 * num_active_params + attention_flops = 12 * num_attention_layers * sequence_length * n_head_q * head_dim + return dense_flops + attention_flops + + @staticmethod + def count_active_parameters(model: torch.nn.Module) -> int: + """ + Counts the parameters a single token actually visits. + + Every parameter counts once, except the routed experts of an MoE layer: only ``top_k`` of + ``num_experts`` are evaluated per token, so their contribution is scaled accordingly. The + router, the shared experts and all dense layers are always active. + + Args: + model (nn.Module): The (unwrapped) model. + + Returns: + int: The number of active parameters. + """ + from modalities.models.components.moe.moe import MoE + + total = sum(p.numel() for p in model.parameters()) + for module in model.modules(): + if not isinstance(module, MoE): + continue + routed = sum(p.numel() for p in module.experts.parameters()) + active_fraction = module.router.top_k / module.router.num_experts + total -= int(routed * (1.0 - active_fraction)) + return total + + def compute(self, num_samples_per_second: torch.Tensor) -> torch.Tensor: + """ + Computes the MFU for a given throughput. + + Args: + num_samples_per_second (torch.Tensor): Observed samples per second. + + Returns: + torch.Tensor: The model FLOPs utilization. + """ + return MFUCalculatorABC._compute_mfu_impl( + num_samples_per_second=num_samples_per_second, + sequence_length=self._sequence_length, + theoretical_flops_per_token=self._theoretical_flops_per_token, + theoretical_gpu_peak_performance=self._theoretical_gpu_peak_performance, + ) diff --git a/tests/fsdp2_parallelization/nemotron_fsdp2_config.yaml b/tests/fsdp2_parallelization/nemotron_fsdp2_config.yaml new file mode 100644 index 000000000..b26521cae --- /dev/null +++ b/tests/fsdp2_parallelization/nemotron_fsdp2_config.yaml @@ -0,0 +1,161 @@ +# Minimal Nemotron hybrid Mamba-Transformer config for distributed FSDP2 tests. +# Exercises all four layer types (Mamba-2, MoE, attention, dense MLP) at a size that fits +# comfortably on two GPUs. + +device_mesh: + component_key: device_mesh + variant_key: default + config: + device_type: cuda + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + world_size: ${cuda_env:WORLD_SIZE} + +initialized_model: + component_key: model + variant_key: model_initialized + config: + model: + instance_key: fsdp_model + pass_type: BY_REFERENCE + model_initializer: + component_key: model_initialization + variant_key: composed + config: + model_type: nemotron + weight_init_type: scaled + mean: 0.0 + std: 0.02 + num_layers: ${model_raw.config.n_layer} + +fsdp_model: + component_key: model + variant_key: fsdp2_wrapped + config: + model: + instance_key: activation_checkpointed_model + pass_type: BY_REFERENCE + device_mesh: + instance_key: device_mesh + pass_type: BY_REFERENCE + mixed_precision_settings: + param_dtype: BF_16 + reduce_dtype: FP_32 + block_names: [Mamba2Layer, NemotronMoELayer, NemotronAttentionLayer, NemotronMLPLayer] + +activation_checkpointed_model: + component_key: model + variant_key: activation_checkpointed + config: + model: + instance_key: model_raw + pass_type: BY_REFERENCE + ac_variant: full_activation_checkpointing + layers_fqn: transformer.h + ac_fun_params: {} + +model_raw: + component_key: model + variant_key: nemotron + config: + use_meta_device: true + sample_key: input_ids + prediction_key: logits + aux_loss_key: moe_aux_loss + sequence_length: 128 + vocab_size: 512 + n_embd: 256 + n_layer: 8 + layer_pattern: "MEMEM*E-" + use_weight_tying: false + lm_head_norm_config: &nemotron_norm_config + norm_type: pytorch_rms_norm + config: + normalized_shape: ${model_raw.config.n_embd} + eps: 1e-5 + layer_specs: + "M": + component_key: nemotron_layer_spec + variant_key: mamba2 + config: + n_embd: ${model_raw.config.n_embd} + mamba_n_heads: 8 + mamba_head_dim: 32 + mamba_state_dim: 16 + mamba_n_groups: 2 + d_conv: 4 + chunk_size: 32 + ssd_backend: native + norm_config: *nemotron_norm_config + "E": + component_key: nemotron_layer_spec + variant_key: moe + config: + n_embd: ${model_raw.config.n_embd} + num_experts: 8 + moe_ffn_hidden: 64 + top_k: 2 + route_scale: 2.5 + score_function: sigmoid + use_expert_bias: true + num_shared_experts: 2 + aux_loss_coeff: 1.0e-4 + experts_backend: grouped_mm + norm_config: *nemotron_norm_config + "*": + component_key: nemotron_layer_spec + variant_key: attention + config: + n_embd: ${model_raw.config.n_embd} + n_head_q: 8 + n_head_kv: 2 + head_dim: 32 + attention_implementation: pytorch_flash + norm_config: *nemotron_norm_config + "-": + component_key: nemotron_layer_spec + variant_key: mlp + config: + n_embd: ${model_raw.config.n_embd} + ffn_hidden: 128 + norm_config: *nemotron_norm_config + +loss_fn: + component_key: loss + variant_key: weighted_sum + config: + weights: [1.0, 1.0] + losses: + - component_key: loss + variant_key: clm_cross_entropy_loss + config: + target_key: target_ids + prediction_key: logits + - component_key: loss + variant_key: moe_aux_loss + config: + prediction_key: moe_aux_loss + +optimizer: + component_key: optimizer + variant_key: moe_load_balanced + config: + expert_bias_update_rate: 1.0e-3 + model: + instance_key: initialized_model + pass_type: BY_REFERENCE + device_mesh: + instance_key: device_mesh + pass_type: BY_REFERENCE + optimizer: + component_key: optimizer + variant_key: adam_w + config: + lr: 1.0e-4 + betas: [0.9, 0.95] + eps: 1e-8 + weight_decay: 0.1 + weight_decay_groups_excluded: [embedding, layernorm, ssm, router] + wrapped_model: + instance_key: initialized_model + pass_type: BY_REFERENCE diff --git a/tests/fsdp2_parallelization/test_nemotron_fsdp2.py b/tests/fsdp2_parallelization/test_nemotron_fsdp2.py new file mode 100644 index 000000000..22113c71d --- /dev/null +++ b/tests/fsdp2_parallelization/test_nemotron_fsdp2.py @@ -0,0 +1,147 @@ +"""Distributed FSDP2 tests for the Nemotron hybrid Mamba-Transformer. + +These cover the parts that only manifest under real sharding: + +* the mixture-of-experts weights are sharded along the expert dimension, +* the state space parameters survive meta-device materialization and initialization, +* a full forward/backward/optimizer step runs, including the auxiliary-loss-free expert bias + update, which reduces token counts across data-parallel ranks. +""" + +import os +from pathlib import Path + +import pytest +import torch +import torch.multiprocessing as mp +from pydantic import BaseModel + +from modalities.__main__ import Main +from modalities.batch import InferenceResultBatch +from modalities.config.config import ProcessGroupBackendType +from modalities.config.pydantic_if_types import ( + PydanticDeviceMeshIFType, + PydanticFSDP2ModuleType, + PydanticLossIFType, + PydanticOptimizerIFType, +) +from tests.end2end_tests.custom_components import MultiProcessingCudaEnv + +CONFIG_PATH = Path(os.path.dirname(__file__)) / "nemotron_fsdp2_config.yaml" +WORLD_SIZE = 2 +VOCAB_SIZE = 512 +SEQ_LEN = 32 + + +class _Components(BaseModel): + initialized_model: PydanticFSDP2ModuleType + device_mesh: PydanticDeviceMeshIFType + optimizer: PydanticOptimizerIFType + loss_fn: PydanticLossIFType + + +def _build_components(tmp_path: Path) -> _Components: + main_obj = Main(CONFIG_PATH, experiments_root_path=tmp_path) + return main_obj.build_components(components_model_type=_Components) + + +def _sharding_worker(process_id: int, tmp_path: str, rdvz_port: int): + with MultiProcessingCudaEnv( + process_group_backend=ProcessGroupBackendType.nccl, + global_rank=process_id, + local_rank=process_id, + world_size=WORLD_SIZE, + rdvz_port=rdvz_port, + ): + components = _build_components(Path(tmp_path)) + model = components.initialized_model + + # Activation checkpointing inserts wrapper modules into the fully qualified names, so match + # on the suffix rather than on an exact path. + parameters = dict(model.named_parameters()) + + def find(suffix: str) -> torch.Tensor: + matches = [param for name, param in parameters.items() if name.endswith(suffix)] + assert matches, f"no parameter ending in {suffix!r}; available: {sorted(parameters)}" + return matches[0] + + # Expert weights are stored as (num_experts, ffn_hidden, n_embd); FSDP2 shards dim 0, so + # each rank owns a slice of the experts. + expert_weight = find("moe.experts.w1") + assert expert_weight.shape == (8, 64, 256), expert_weight.shape + local_expert_weight = expert_weight.to_local() + assert local_expert_weight.shape[0] == 8 // WORLD_SIZE, local_expert_weight.shape + + # Every parameter must be materialized and finite after initialization. + for name, param in parameters.items(): + local = param.to_local() if hasattr(param, "to_local") else param + assert local.device.type != "meta", f"{name} is still on the meta device" + if local.numel() > 0: + assert torch.isfinite(local).all(), f"{name} is not finite" + + # The state space parameters must have kept their own distributions, not the normal one. + A_log = find("mixer.A_log") + A = torch.exp(A_log.full_tensor() if hasattr(A_log, "full_tensor") else A_log) + assert A.min() >= 1.0 and A.max() <= 16.0, (A.min().item(), A.max().item()) + + D = find("mixer.D") + D_full = D.full_tensor() if hasattr(D, "full_tensor") else D + torch.testing.assert_close(D_full, torch.ones_like(D_full)) + + # The router gate, by contrast, must have been touched by the normal-distribution + # initializer, confirming the two initialization paths coexist correctly. + gate = find("moe.router.gate.weight") + gate_full = gate.full_tensor() if hasattr(gate, "full_tensor") else gate + assert gate_full.std().item() == pytest.approx(0.02, rel=0.3) + + +def _training_step_worker(process_id: int, tmp_path: str, rdvz_port: int): + with MultiProcessingCudaEnv( + process_group_backend=ProcessGroupBackendType.nccl, + global_rank=process_id, + local_rank=process_id, + world_size=WORLD_SIZE, + rdvz_port=rdvz_port, + ): + components = _build_components(Path(tmp_path)) + model = components.initialized_model + optimizer = components.optimizer + loss_fn = components.loss_fn + + # Different data per rank, so that the expert-load reduction has something to combine. + generator = torch.Generator(device="cuda").manual_seed(process_id) + inputs = torch.randint(0, VOCAB_SIZE, (2, SEQ_LEN), device="cuda", generator=generator) + targets = torch.randint(0, VOCAB_SIZE, (2, SEQ_LEN), device="cuda", generator=generator) + + losses = [] + for _ in range(2): + predictions = model({"input_ids": inputs}) + assert "moe_aux_loss" in predictions + batch = InferenceResultBatch(targets={"target_ids": targets}, predictions=predictions) + loss = loss_fn(batch) + assert torch.isfinite(loss), loss + loss.backward() + optimizer.step() + optimizer.zero_grad(set_to_none=True) + losses.append(loss.item()) + + # The load-balancing hook must have moved the expert biases away from zero and reset the + # token counters. Both are rank-consistent because the counts are all-reduced. + moe_layers = [module for name, module in model.named_modules() if name.endswith(".moe")] + assert moe_layers, "no MoE layers found" + for moe in moe_layers: + expert_bias = moe.router.expert_bias + assert expert_bias.abs().max() > 0, "expert bias was never updated" + assert moe.router.tokens_per_expert.sum() == 0, "token counters were not reset" + + assert all(loss == loss for loss in losses) # no NaNs + + +@pytest.mark.skipif(torch.cuda.device_count() < WORLD_SIZE, reason=f"This test requires {WORLD_SIZE} GPUs") +def test_nemotron_fsdp2_shards_experts_and_initializes_state_space_parameters(tmp_path): + mp.spawn(_sharding_worker, args=(str(tmp_path), 22421), nprocs=WORLD_SIZE, join=True) + + +@pytest.mark.skipif(torch.cuda.device_count() < WORLD_SIZE, reason=f"This test requires {WORLD_SIZE} GPUs") +def test_nemotron_fsdp2_training_step_with_load_balancing(tmp_path): + mp.spawn(_training_step_worker, args=(str(tmp_path), 22422), nprocs=WORLD_SIZE, join=True) diff --git a/tests/models/components/test_norms.py b/tests/models/components/test_norms.py new file mode 100644 index 000000000..0b87a5167 --- /dev/null +++ b/tests/models/components/test_norms.py @@ -0,0 +1,39 @@ +import torch.nn as nn + +from modalities.models.components.layer_norms import RMSLayerNorm +from modalities.models.components.norms import NormWrapperConfig + + +def test_norm_wrapper_builds_pytorch_rms_norm(): + config = NormWrapperConfig.model_validate( + {"norm_type": "pytorch_rms_norm", "config": {"normalized_shape": 16, "eps": 1e-5}} + ) + norm = config.build() + assert isinstance(norm, nn.RMSNorm) + assert norm.normalized_shape == (16,) + assert norm.eps == 1e-5 + + +def test_norm_wrapper_builds_layer_norm(): + config = NormWrapperConfig.model_validate( + {"norm_type": "layer_norm", "config": {"normalized_shape": 8, "eps": 1e-6, "bias": False}} + ) + norm = config.build() + assert isinstance(norm, nn.LayerNorm) + assert norm.bias is None + + +def test_norm_wrapper_builds_deprecated_rms_norm(): + config = NormWrapperConfig.model_validate({"norm_type": "rms_norm", "config": {"ndim": 8}}) + assert isinstance(config.build(), RMSLayerNorm) + + +def test_norm_wrapper_build_returns_a_fresh_instance_per_call(): + # Layers must not share normalization parameters, so build() has to create a new module + # every time it is called. + config = NormWrapperConfig.model_validate( + {"norm_type": "pytorch_rms_norm", "config": {"normalized_shape": 4, "eps": 1e-5}} + ) + first, second = config.build(), config.build() + assert first is not second + assert first.weight is not second.weight diff --git a/tests/models/nemotron/__init__.py b/tests/models/nemotron/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/models/nemotron/test_layer_pattern.py b/tests/models/nemotron/test_layer_pattern.py new file mode 100644 index 000000000..19b49dce6 --- /dev/null +++ b/tests/models/nemotron/test_layer_pattern.py @@ -0,0 +1,56 @@ +import pytest + +from modalities.models.nemotron.layer_pattern import ( + LayerSymbol, + count_layers_by_type, + get_num_layers, + parse_layer_pattern, +) + +# The published Nemotron-3 Nano 30B-A3B pattern (52 layers). +NEMOTRON_3_NANO_PATTERN = "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME" + + +def test_parse_layer_pattern_returns_one_symbol_per_character(): + assert parse_layer_pattern("ME*-") == [ + LayerSymbol.MAMBA, + LayerSymbol.MOE, + LayerSymbol.ATTENTION, + LayerSymbol.MLP, + ] + + +def test_parse_layer_pattern_rejects_empty_pattern(): + with pytest.raises(ValueError, match="must not be empty"): + parse_layer_pattern("") + + +@pytest.mark.parametrize("pattern", ["MEX", "M|M", "M/M", "m"]) +def test_parse_layer_pattern_rejects_unknown_symbols(pattern): + with pytest.raises(ValueError, match="Invalid layer symbol"): + parse_layer_pattern(pattern) + + +def test_parse_layer_pattern_error_reports_position(): + with pytest.raises(ValueError, match="at position 2"): + parse_layer_pattern("MEX") + + +def test_count_layers_by_type_includes_absent_types(): + counts = count_layers_by_type("MME") + assert counts == { + LayerSymbol.MAMBA: 2, + LayerSymbol.MOE: 1, + LayerSymbol.ATTENTION: 0, + LayerSymbol.MLP: 0, + } + + +def test_nemotron_3_nano_pattern_matches_published_architecture(): + # Model report Table 1 / Figure 2: 52 layers of which 6 are self-attention. + assert get_num_layers(NEMOTRON_3_NANO_PATTERN) == 52 + counts = count_layers_by_type(NEMOTRON_3_NANO_PATTERN) + assert counts[LayerSymbol.ATTENTION] == 6 + assert counts[LayerSymbol.MAMBA] == 23 + assert counts[LayerSymbol.MOE] == 23 + assert counts[LayerSymbol.MLP] == 0 diff --git a/tests/models/nemotron/test_mamba2_mixer.py b/tests/models/nemotron/test_mamba2_mixer.py new file mode 100644 index 000000000..0c8bff9b9 --- /dev/null +++ b/tests/models/nemotron/test_mamba2_mixer.py @@ -0,0 +1,236 @@ +import importlib.util +import math + +import pytest +import torch + +from modalities.models.components.mamba2.mamba2_mixer import Mamba2Mixer, SSDBackend + +# The fused Triton kernels come from the optional `mamba` extra. +HAS_FUSED_KERNELS = all( + importlib.util.find_spec(module_name) is not None for module_name in ("mamba_ssm", "causal_conv1d") +) + +# Small but structurally faithful configuration: n_heads is a multiple of n_groups and the inner +# dimension (n_heads * head_dim) differs from n_embd, as in the real Nemotron-3 Nano. +MIXER_KWARGS = dict( + n_embd=32, + n_heads=4, + head_dim=8, + state_dim=6, + n_groups=2, + d_conv=4, + chunk_size=4, +) + + +def _make_mixer(**overrides) -> Mamba2Mixer: + torch.manual_seed(0) + return Mamba2Mixer(**{**MIXER_KWARGS, **overrides}) + + +def test_mixer_parameter_shapes_match_reference_layout(): + mixer = _make_mixer() + d_inner = MIXER_KWARGS["n_heads"] * MIXER_KWARGS["head_dim"] + group_state = MIXER_KWARGS["n_groups"] * MIXER_KWARGS["state_dim"] + conv_dim = d_inner + 2 * group_state + + # in_proj packs [z, x, B, C, dt]. + assert mixer.in_proj.weight.shape == (2 * d_inner + 2 * group_state + MIXER_KWARGS["n_heads"], 32) + assert mixer.in_proj.bias is None + assert mixer.conv1d_weight.shape == (conv_dim, 1, MIXER_KWARGS["d_conv"]) + assert mixer.conv1d_bias.shape == (conv_dim,) + assert mixer.A_log.shape == (MIXER_KWARGS["n_heads"],) + assert mixer.D.shape == (MIXER_KWARGS["n_heads"],) + assert mixer.dt_bias.shape == (MIXER_KWARGS["n_heads"],) + assert mixer.out_proj.weight.shape == (32, d_inner) + assert mixer.d_inner == d_inner + assert mixer.conv_dim == conv_dim + + +def test_nemotron_3_nano_mixer_dimensions(): + # Model report Table 1: model dim 2688, 64 Mamba heads of dim 64, state dim 128, 8 groups. + mixer = Mamba2Mixer(n_embd=2688, n_heads=64, head_dim=64, state_dim=128, n_groups=8, chunk_size=128) + assert mixer.d_inner == 4096 + assert mixer.in_proj.out_features == 2 * 4096 + 2 * 8 * 128 + 64 == 10304 + assert mixer.conv_dim == 4096 + 2 * 8 * 128 == 6144 + assert mixer.norm.group_size == 512 + + +def test_mixer_forward_preserves_shape(): + mixer = _make_mixer() + x = torch.randn(2, 12, MIXER_KWARGS["n_embd"]) + out = mixer(x) + assert out.shape == x.shape + assert torch.isfinite(out).all() + + +@pytest.mark.parametrize("seq_len", [1, 3, 8, 13]) +def test_mixer_handles_arbitrary_sequence_lengths(seq_len): + mixer = _make_mixer() + out = mixer(torch.randn(2, seq_len, MIXER_KWARGS["n_embd"])) + assert out.shape == (2, seq_len, MIXER_KWARGS["n_embd"]) + assert torch.isfinite(out).all() + + +def test_mixer_is_causal(): + # The defining property of a sequence mixer used for causal LM: perturbing token t must not + # change any output at a position before t. This covers the conv, the scan and the gating. + mixer = _make_mixer().eval() + x = torch.randn(1, 16, MIXER_KWARGS["n_embd"]) + with torch.no_grad(): + baseline = mixer(x) + perturbed = x.clone() + perturbed[:, 10] += 5.0 + outputs = mixer(perturbed) + + torch.testing.assert_close(outputs[:, :10], baseline[:, :10], rtol=1e-4, atol=1e-5) + assert not torch.allclose(outputs[:, 10], baseline[:, 10]) + + +@pytest.mark.parametrize("chunk_size", [1, 2, 4, 16]) +def test_mixer_output_is_independent_of_chunk_size(chunk_size): + reference = _make_mixer(chunk_size=16).eval() + candidate = _make_mixer(chunk_size=chunk_size).eval() + candidate.load_state_dict(reference.state_dict()) + + x = torch.randn(2, 16, MIXER_KWARGS["n_embd"]) + with torch.no_grad(): + torch.testing.assert_close(candidate(x), reference(x), rtol=1e-4, atol=1e-4) + + +def test_mixer_batch_elements_are_independent(): + mixer = _make_mixer().eval() + x = torch.randn(3, 10, MIXER_KWARGS["n_embd"]) + with torch.no_grad(): + batched = mixer(x) + individually = torch.cat([mixer(x[i : i + 1]) for i in range(3)], dim=0) + torch.testing.assert_close(batched, individually, rtol=1e-4, atol=1e-5) + + +def test_mixer_is_differentiable(): + mixer = _make_mixer() + x = torch.randn(2, 8, MIXER_KWARGS["n_embd"], requires_grad=True) + mixer(x).sum().backward() + + assert x.grad is not None and torch.isfinite(x.grad).all() + for name, param in mixer.named_parameters(): + assert param.grad is not None, f"{name} received no gradient" + assert torch.isfinite(param.grad).all(), f"{name} has non-finite gradient" + + +def test_reset_parameters_produces_the_reference_distributions(): + mixer = _make_mixer(n_heads=64, head_dim=8, n_groups=8) + mixer.reset_parameters() + + # A = -exp(A_log) must lie in the sampled range [1, 16] in magnitude. + A = torch.exp(mixer.A_log) + assert A.min() >= 1.0 and A.max() <= 16.0 + assert mixer.A_log.dtype == torch.float32 + + # softplus(dt_bias) must reproduce dt in [dt_min, dt_max]. + dt = torch.nn.functional.softplus(mixer.dt_bias) + assert dt.min() >= mixer.dt_min * 0.99 + assert dt.max() <= mixer.dt_max * 1.01 + + # D is the identity skip. + torch.testing.assert_close(mixer.D, torch.ones_like(mixer.D)) + + # conv1d uses the default nn.Conv1d bounds. + bound = 1.0 / math.sqrt(mixer.conv1d_weight.size(1) * mixer.conv1d_weight.size(2)) + assert mixer.conv1d_bias.abs().max() <= bound + + +def test_reset_parameters_is_a_noop_on_meta_device(): + # The model factory builds on meta and materializes later; reset_parameters must not raise. + with torch.device("meta"): + mixer = Mamba2Mixer(**MIXER_KWARGS) + mixer.reset_parameters() + assert mixer.A_log.is_meta + + +@pytest.mark.skipif(HAS_FUSED_KERNELS, reason="mamba-ssm and causal-conv1d are installed") +def test_fused_backend_raises_a_helpful_error_when_kernels_are_missing(): + # Requesting a backend whose dependencies are absent must fail loudly at construction time + # rather than silently falling back to a slower path. + with pytest.raises(ValueError, match="requires mamba-ssm and causal-conv1d"): + _make_mixer(ssd_backend=SSDBackend.FUSED) + + +def test_mixer_rejects_inconsistent_head_and_group_counts(): + with pytest.raises(ValueError, match="must be divisible by n_groups"): + _make_mixer(n_heads=3, n_groups=2) + + +def test_mixer_rejects_invalid_kernel_and_init_ranges(): + with pytest.raises(ValueError, match="d_conv must be at least 1"): + _make_mixer(d_conv=0) + with pytest.raises(ValueError, match="A_init_range"): + _make_mixer(A_init_range=(0.0, 16.0)) + with pytest.raises(ValueError, match="A_init_range"): + _make_mixer(A_init_range=(16.0, 1.0)) + + +def test_mixer_warns_when_native_backend_is_used_at_scale(caplog): + with caplog.at_level("WARNING"): + Mamba2Mixer(n_embd=2048, n_heads=8, head_dim=8, state_dim=4, n_groups=2) + assert "native (pure-PyTorch) SSD backend" in caplog.text + + +@pytest.mark.parametrize("ssd_backend", [SSDBackend.NATIVE, "native"]) +def test_backend_accepts_enum_and_string(ssd_backend): + mixer = _make_mixer(ssd_backend=ssd_backend) + assert mixer.ssd_backend == SSDBackend.NATIVE + + +# ------------------------------------------------------------------------------------------------ +# Fused backend parity. Skipped unless the optional `mamba` extra is installed and a GPU is present. +# ------------------------------------------------------------------------------------------------ + +requires_fused_kernels = pytest.mark.skipif( + not (HAS_FUSED_KERNELS and torch.cuda.is_available()), + reason="requires the optional `mamba` extra and a CUDA device", +) + +# The fused kernels require the head dimension to be a multiple of 8 and prefer realistic sizes. +FUSED_MIXER_KWARGS = dict(n_embd=256, n_heads=8, head_dim=32, state_dim=32, n_groups=2, d_conv=4, chunk_size=64) + + +@requires_fused_kernels +def test_fused_backend_matches_native_backend(): + # The native scan is validated against a step-by-step recurrence elsewhere, so agreement here + # transitively validates the fused path against the mathematical definition. + torch.manual_seed(0) + native = Mamba2Mixer(**FUSED_MIXER_KWARGS, ssd_backend=SSDBackend.NATIVE).cuda().bfloat16().eval() + fused = Mamba2Mixer(**FUSED_MIXER_KWARGS, ssd_backend=SSDBackend.FUSED).cuda().bfloat16().eval() + fused.load_state_dict(native.state_dict()) + + x = torch.randn(2, 128, FUSED_MIXER_KWARGS["n_embd"], device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + torch.testing.assert_close(fused(x), native(x), rtol=3e-2, atol=3e-2) + + +@requires_fused_kernels +def test_fused_backend_is_causal(): + mixer = Mamba2Mixer(**FUSED_MIXER_KWARGS, ssd_backend=SSDBackend.FUSED).cuda().bfloat16().eval() + x = torch.randn(1, 128, FUSED_MIXER_KWARGS["n_embd"], device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + baseline = mixer(x) + perturbed = x.clone() + perturbed[:, 70] += 5.0 + outputs = mixer(perturbed) + torch.testing.assert_close(outputs[:, :70], baseline[:, :70], rtol=2e-2, atol=2e-2) + + +@requires_fused_kernels +def test_fused_backend_gradients_match_native_backend(): + torch.manual_seed(0) + native = Mamba2Mixer(**FUSED_MIXER_KWARGS, ssd_backend=SSDBackend.NATIVE).cuda().bfloat16() + fused = Mamba2Mixer(**FUSED_MIXER_KWARGS, ssd_backend=SSDBackend.FUSED).cuda().bfloat16() + fused.load_state_dict(native.state_dict()) + + x = torch.randn(2, 128, FUSED_MIXER_KWARGS["n_embd"], device="cuda", dtype=torch.bfloat16) + native_in, fused_in = x.clone().requires_grad_(True), x.clone().requires_grad_(True) + native(native_in).float().sum().backward() + fused(fused_in).float().sum().backward() + torch.testing.assert_close(fused_in.grad, native_in.grad, rtol=5e-2, atol=5e-2) diff --git a/tests/models/nemotron/test_moe.py b/tests/models/nemotron/test_moe.py new file mode 100644 index 000000000..eec31cb8a --- /dev/null +++ b/tests/models/nemotron/test_moe.py @@ -0,0 +1,440 @@ +import pytest +import torch + +from modalities.models.components.moe.experts import ExpertsBackend, GroupedExperts, squared_relu +from modalities.models.components.moe.moe import MoE +from modalities.models.components.moe.router import RouterScoreFunction, TopKRouter +from modalities.models.nemotron.nemotron_mlp import SquaredReLUMLP + +N_EMBD = 16 +NUM_EXPERTS = 8 +TOP_K = 2 +# Both matmul inner dimensions must be multiples of 8 so that the grouped_mm backend is usable. +FFN_HIDDEN = 24 + + +def _make_router(**overrides) -> TopKRouter: + torch.manual_seed(0) + kwargs = dict(n_embd=N_EMBD, num_experts=NUM_EXPERTS, top_k=TOP_K) + return TopKRouter(**{**kwargs, **overrides}) + + +def _make_experts(**overrides) -> GroupedExperts: + torch.manual_seed(0) + kwargs = dict(n_embd=N_EMBD, ffn_hidden=FFN_HIDDEN, num_experts=NUM_EXPERTS, backend=ExpertsBackend.LOOPED) + experts = GroupedExperts(**{**kwargs, **overrides}) + with torch.no_grad(): + experts.w1.normal_(0, 0.02) + experts.w2.normal_(0, 0.02) + return experts + + +def _make_moe(aux_loss_coeff: float = 0.0, with_shared: bool = False) -> MoE: + shared = SquaredReLUMLP(n_embd=N_EMBD, ffn_hidden=2 * FFN_HIDDEN) if with_shared else None + return MoE( + router=_make_router(), + experts=_make_experts(), + shared_experts=shared, + aux_loss_coeff=aux_loss_coeff, + ) + + +# -------------------------------------------------------------------------------------------- +# Activation +# -------------------------------------------------------------------------------------------- + + +def test_squared_relu_zeroes_negatives_and_squares_positives(): + x = torch.tensor([-2.0, -0.5, 0.0, 0.5, 3.0]) + torch.testing.assert_close(squared_relu(x), torch.tensor([0.0, 0.0, 0.0, 0.25, 9.0])) + + +# -------------------------------------------------------------------------------------------- +# Router +# -------------------------------------------------------------------------------------------- + + +def test_router_output_shapes_and_dtypes(): + router = _make_router() + x = torch.randn(10, N_EMBD) + weights, indices, scores = router(x) + + assert weights.shape == (10, TOP_K) + assert indices.shape == (10, TOP_K) + assert scores.shape == (10, NUM_EXPERTS) + assert weights.dtype == x.dtype + assert indices.dtype == torch.int64 + # Scores stay in the router dtype so the auxiliary loss is computed in fp32. + assert scores.dtype == torch.float32 + + +def test_router_selects_distinct_experts_per_token(): + router = _make_router() + _, indices, _ = router(torch.randn(32, N_EMBD)) + for row in indices: + assert len(set(row.tolist())) == TOP_K + + +def test_sigmoid_router_weights_sum_to_route_scale(): + router = _make_router(route_scale=2.5) + weights, _, _ = router(torch.randn(20, N_EMBD)) + # Renormalization over the selected experts followed by the constant rescaling. + torch.testing.assert_close(weights.sum(dim=-1), torch.full((20,), 2.5), rtol=1e-5, atol=1e-5) + + +def test_sigmoid_router_with_top_k_one_does_not_renormalize(): + # With top_k == 1 the reference implementation keeps the raw sigmoid score rather than + # normalizing it to 1, so the weight must stay below 1. + router = _make_router(top_k=1, route_scale=1.0) + weights, indices, scores = router(torch.randn(20, N_EMBD)) + torch.testing.assert_close(weights.squeeze(-1), torch.gather(scores, 1, indices).squeeze(-1)) + assert (weights < 1.0).all() + + +def test_softmax_router_weights_form_a_distribution(): + router = _make_router(score_function=RouterScoreFunction.SOFTMAX) + weights, _, scores = router(torch.randn(20, N_EMBD)) + torch.testing.assert_close(weights.sum(dim=-1), torch.ones(20), rtol=1e-5, atol=1e-5) + torch.testing.assert_close(scores.sum(dim=-1), torch.ones(20), rtol=1e-5, atol=1e-5) + + +def test_expert_bias_changes_selection_but_not_returned_weights(): + # The load-balancing bias must steer which experts are picked without leaking into the + # forward signal; otherwise balancing would distort the model's output. + router = _make_router(route_scale=1.0) + x = torch.randn(64, N_EMBD) + _, baseline_indices, scores = router(x) + + starved_expert = 5 + with torch.no_grad(): + router.expert_bias[starved_expert] = 100.0 + weights, indices, _ = router(x) + + assert (indices == starved_expert).any() + assert not torch.equal(indices, baseline_indices) + # Weights are still the unbiased sigmoid scores of the selected experts (renormalized). + selected = torch.gather(scores, 1, indices) + torch.testing.assert_close(weights, selected / selected.sum(dim=-1, keepdim=True), rtol=1e-5, atol=1e-6) + + +def test_router_without_expert_bias_has_no_bias_buffer(): + router = _make_router(use_expert_bias=False) + assert router.expert_bias is None + assert "expert_bias" not in dict(router.named_buffers()) + + +def test_expert_bias_is_persistent_and_token_counter_is_not(): + router = _make_router() + state_dict = router.state_dict() + assert "expert_bias" in state_dict + assert "tokens_per_expert" not in state_dict + + +def test_router_rejects_invalid_top_k(): + with pytest.raises(ValueError, match=r"top_k must be in"): + _make_router(top_k=0) + with pytest.raises(ValueError, match=r"top_k must be in"): + _make_router(top_k=NUM_EXPERTS + 1) + + +def test_router_reset_parameters_zeroes_load_balancing_state(): + router = _make_router() + with torch.no_grad(): + router.expert_bias.fill_(3.0) + router.tokens_per_expert.fill_(7.0) + router.reset_parameters() + torch.testing.assert_close(router.expert_bias, torch.zeros(NUM_EXPERTS)) + torch.testing.assert_close(router.tokens_per_expert, torch.zeros(NUM_EXPERTS)) + + +# -------------------------------------------------------------------------------------------- +# Grouped experts +# -------------------------------------------------------------------------------------------- + + +def test_grouped_experts_parameter_shapes_are_non_gated(): + experts = _make_experts() + # Two matrices per expert (no gate projection), expert dimension first for FSDP sharding. + assert experts.w1.shape == (NUM_EXPERTS, FFN_HIDDEN, N_EMBD) + assert experts.w2.shape == (NUM_EXPERTS, N_EMBD, FFN_HIDDEN) + assert len(list(experts.parameters())) == 2 + + +def test_looped_experts_apply_the_right_expert_to_each_slice(): + experts = _make_experts() + tokens_per_expert = torch.tensor([2, 0, 3, 1, 0, 0, 0, 0]) + x_sorted = torch.randn(int(tokens_per_expert.sum()), N_EMBD) + out = experts(x_sorted, tokens_per_expert) + + start = 0 + for expert_idx, count in enumerate(tokens_per_expert.tolist()): + if count == 0: + continue + chunk = x_sorted[start : start + count] + expected = squared_relu(chunk @ experts.w1[expert_idx].T) @ experts.w2[expert_idx].T + torch.testing.assert_close(out[start : start + count], expected, rtol=1e-5, atol=1e-6) + start += count + + +def test_experts_handle_all_tokens_on_one_expert(): + experts = _make_experts() + tokens_per_expert = torch.zeros(NUM_EXPERTS, dtype=torch.int64) + tokens_per_expert[3] = 5 + out = experts(torch.randn(5, N_EMBD), tokens_per_expert) + assert out.shape == (5, N_EMBD) + assert torch.isfinite(out).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="torch._grouped_mm requires CUDA") +def test_grouped_mm_backend_matches_looped_backend(): + torch.manual_seed(0) + looped = _make_experts(backend=ExpertsBackend.LOOPED).cuda().bfloat16() + grouped = _make_experts(backend=ExpertsBackend.GROUPED_MM).cuda().bfloat16() + grouped.load_state_dict(looped.state_dict()) + + tokens_per_expert = torch.tensor([4, 2, 0, 6, 1, 3, 0, 8], device="cuda") + x_sorted = torch.randn(int(tokens_per_expert.sum()), N_EMBD, device="cuda", dtype=torch.bfloat16) + torch.testing.assert_close( + grouped(x_sorted, tokens_per_expert), looped(x_sorted, tokens_per_expert), rtol=2e-2, atol=2e-2 + ) + + +def test_grouped_mm_backend_falls_back_on_cpu(): + # The backend selection must degrade gracefully so that CPU tests and CPU inference work. + experts = _make_experts(backend=ExpertsBackend.GROUPED_MM) + tokens_per_expert = torch.tensor([2, 1, 0, 0, 0, 0, 0, 1]) + out = experts(torch.randn(4, N_EMBD), tokens_per_expert) + assert out.shape == (4, N_EMBD) + assert torch.isfinite(out).all() + + +@pytest.mark.parametrize("n_embd,ffn_hidden", [(12, 24), (16, 20), (12, 12)]) +def test_grouped_mm_backend_rejects_unaligned_dimensions(n_embd, ffn_hidden): + # torch._grouped_mm needs 16-byte aligned strides. Catching that at construction time is far + # more useful than a "strides should be multiple of 16 bytes" RuntimeError mid-training. + with pytest.raises(ValueError, match="must be a multiple of 8"): + GroupedExperts(n_embd=n_embd, ffn_hidden=ffn_hidden, num_experts=NUM_EXPERTS, backend=ExpertsBackend.GROUPED_MM) + + +def test_looped_backend_accepts_unaligned_dimensions(): + experts = GroupedExperts(n_embd=12, ffn_hidden=12, num_experts=2, backend=ExpertsBackend.LOOPED) + with torch.no_grad(): + experts.w1.normal_(0, 0.02) + experts.w2.normal_(0, 0.02) + out = experts(torch.randn(3, 12), torch.tensor([1, 2])) + assert out.shape == (3, 12) + + +# -------------------------------------------------------------------------------------------- +# MoE layer +# -------------------------------------------------------------------------------------------- + + +def test_moe_forward_preserves_shape(): + moe = _make_moe() + x = torch.randn(2, 6, N_EMBD) + out = moe(x) + assert out.shape == x.shape + assert torch.isfinite(out).all() + + +def test_moe_matches_an_explicit_per_token_reference(): + # The dispatch/sort/combine machinery is the easiest place to introduce an indexing bug, so + # compare against a naive loop over tokens. + moe = _make_moe() + x = torch.randn(2, 5, N_EMBD) + out = moe(x) + + x_flat = x.reshape(-1, N_EMBD) + weights, indices, _ = moe.router(x_flat) + expected = torch.zeros_like(x_flat) + for token_idx in range(x_flat.shape[0]): + for slot in range(TOP_K): + expert_idx = indices[token_idx, slot] + hidden = squared_relu(x_flat[token_idx] @ moe.experts.w1[expert_idx].T) + expected[token_idx] += weights[token_idx, slot] * (hidden @ moe.experts.w2[expert_idx].T) + + torch.testing.assert_close(out.reshape(-1, N_EMBD), expected, rtol=1e-4, atol=1e-5) + + +def test_moe_shared_expert_output_is_added(): + routed_only = _make_moe(with_shared=False) + with_shared = _make_moe(with_shared=True) + with_shared.router.load_state_dict(routed_only.router.state_dict()) + with_shared.experts.load_state_dict(routed_only.experts.state_dict()) + + x = torch.randn(2, 4, N_EMBD) + expected = routed_only(x) + with_shared.shared_experts(x) + torch.testing.assert_close(with_shared(x), expected, rtol=1e-5, atol=1e-6) + + +def test_moe_counts_tokens_per_expert(): + moe = _make_moe() + x = torch.randn(2, 8, N_EMBD) + moe(x) + # Every token contributes exactly top_k routed slots. + assert moe.router.tokens_per_expert.sum().item() == 2 * 8 * TOP_K + + +def test_moe_token_counts_accumulate_across_micro_batches(): + # Gradient accumulation means several forwards per optimizer step; the counter must sum them. + moe = _make_moe() + moe(torch.randn(1, 4, N_EMBD)) + after_first = moe.router.tokens_per_expert.clone() + moe(torch.randn(1, 4, N_EMBD)) + assert moe.router.tokens_per_expert.sum() == 2 * after_first.sum() + + +def test_moe_aux_loss_is_none_when_disabled(): + moe = _make_moe(aux_loss_coeff=0.0) + moe(torch.randn(2, 4, N_EMBD)) + assert moe.last_aux_loss is None + + +def test_moe_aux_loss_is_a_scalar_in_the_autograd_graph(): + moe = _make_moe(aux_loss_coeff=1e-2) + moe(torch.randn(2, 4, N_EMBD)) + assert moe.last_aux_loss.ndim == 0 + assert moe.last_aux_loss.requires_grad + moe.last_aux_loss.backward() + assert moe.router.gate.weight.grad is not None + + +def test_moe_aux_loss_is_overwritten_not_accumulated(): + # Overwriting is what makes the loss correct under activation checkpointing, where the + # forward pass is replayed during the backward pass. + moe = _make_moe(aux_loss_coeff=1e-2) + x = torch.randn(2, 4, N_EMBD) + moe(x) + first = moe.last_aux_loss.item() + moe(x) + assert moe.last_aux_loss.item() == pytest.approx(first) + + +def test_moe_aux_loss_is_minimal_for_perfectly_balanced_routing(): + # The Switch loss is E * with sum(f) == sum(P) == 1. It attains its floor of 1.0 when + # both the load f and the router probability P are uniform, and grows as the two become + # correlated (i.e. the router keeps preferring the experts that are already overloaded). + moe = _make_moe(aux_loss_coeff=1.0) + num_tokens = 64 + uniform_scores = torch.full((num_tokens, NUM_EXPERTS), 1.0 / NUM_EXPERTS) + # Round-robin assignment gives every expert exactly the same number of slots. + balanced_indices = torch.stack( + [torch.arange(num_tokens) % NUM_EXPERTS, (torch.arange(num_tokens) + 1) % NUM_EXPERTS], dim=-1 + ) + balanced = moe._compute_aux_loss(scores=uniform_scores, top_indices=balanced_indices, batch_size=1) + torch.testing.assert_close(balanced, torch.tensor(1.0), rtol=1e-5, atol=1e-6) + + # Collapse: the router puts all its probability mass on experts 0 and 1 and also routes + # every token there, which is the pathology the loss is meant to punish. + peaked_scores = torch.full((num_tokens, NUM_EXPERTS), 1e-6) + peaked_scores[:, 0] = 0.5 + peaked_scores[:, 1] = 0.5 + collapsed_indices = torch.zeros_like(balanced_indices) + collapsed_indices[:, 1] = 1 + collapsed = moe._compute_aux_loss(scores=peaked_scores, top_indices=collapsed_indices, batch_size=1) + assert collapsed > balanced + + +def test_moe_aux_loss_is_invariant_to_the_load_when_probabilities_are_uniform(): + # A property worth pinning down explicitly: with uniform P the loss cannot distinguish loads, + # because sum(f) is always 1. The gradient therefore acts on the router probabilities, not on + # the (non-differentiable) assignment counts. + moe = _make_moe(aux_loss_coeff=1.0) + num_tokens = 32 + uniform_scores = torch.full((num_tokens, NUM_EXPERTS), 1.0 / NUM_EXPERTS) + balanced_indices = torch.stack( + [torch.arange(num_tokens) % NUM_EXPERTS, (torch.arange(num_tokens) + 1) % NUM_EXPERTS], dim=-1 + ) + collapsed_indices = torch.zeros_like(balanced_indices) + collapsed_indices[:, 1] = 1 + + balanced = moe._compute_aux_loss(scores=uniform_scores, top_indices=balanced_indices, batch_size=1) + collapsed = moe._compute_aux_loss(scores=uniform_scores, top_indices=collapsed_indices, batch_size=1) + torch.testing.assert_close(balanced, collapsed, rtol=1e-6, atol=1e-7) + + +def test_moe_aux_loss_scales_with_the_coefficient(): + x = torch.randn(2, 4, N_EMBD) + small, large = _make_moe(aux_loss_coeff=1e-3), _make_moe(aux_loss_coeff=1e-2) + small(x) + large(x) + torch.testing.assert_close(large.last_aux_loss, small.last_aux_loss * 10.0, rtol=1e-4, atol=1e-8) + + +def test_moe_aux_loss_is_computed_per_sequence(): + # Each sequence collapses onto its own pair of experts. Pooled over the batch this looks + # balanced (four experts used evenly), so a batch-level loss barely reacts. The per-sequence + # loss sees each collapse in isolation and penalizes it. That difference is precisely why + # Nemotron uses the sequence-level variant. + moe = _make_moe(aux_loss_coeff=1.0) + num_tokens_per_seq = 16 + scores = torch.full((2 * num_tokens_per_seq, NUM_EXPERTS), 1e-6) + scores[:num_tokens_per_seq, 0] = 0.5 + scores[:num_tokens_per_seq, 1] = 0.5 + scores[num_tokens_per_seq:, 2] = 0.5 + scores[num_tokens_per_seq:, 3] = 0.5 + + indices = torch.zeros(2 * num_tokens_per_seq, TOP_K, dtype=torch.int64) + indices[:num_tokens_per_seq] = torch.tensor([0, 1]) + indices[num_tokens_per_seq:] = torch.tensor([2, 3]) + + per_sequence = moe._compute_aux_loss(scores=scores, top_indices=indices, batch_size=2) + pooled = moe._compute_aux_loss(scores=scores, top_indices=indices, batch_size=1) + assert per_sequence > pooled + # Sanity: two experts out of eight carrying a whole sequence is a 4x imbalance. + torch.testing.assert_close(per_sequence, torch.tensor(4.0), rtol=1e-3, atol=1e-3) + torch.testing.assert_close(pooled, torch.tensor(2.0), rtol=1e-3, atol=1e-3) + + +def test_moe_is_differentiable_through_router_and_experts(): + moe = _make_moe(with_shared=True) + x = torch.randn(2, 6, N_EMBD, requires_grad=True) + moe(x).sum().backward() + + assert x.grad is not None and torch.isfinite(x.grad).all() + assert moe.router.gate.weight.grad is not None + assert moe.experts.w1.grad is not None + assert moe.experts.w2.grad is not None + assert moe.shared_experts.c_fc.weight.grad is not None + + +def test_moe_gradients_only_reach_selected_experts(): + # Sparsity is the whole point: an expert that received no token must receive no gradient. + moe = _make_moe() + x = torch.randn(1, 2, N_EMBD) + _, indices, _ = moe.router(x.reshape(-1, N_EMBD)) + selected = set(indices.reshape(-1).tolist()) + unselected = sorted(set(range(NUM_EXPERTS)) - selected) + assert unselected, "test requires at least one unrouted expert" + + moe(x).sum().backward() + for expert_idx in unselected: + torch.testing.assert_close(moe.experts.w1.grad[expert_idx], torch.zeros_like(moe.experts.w1[expert_idx])) + + +def test_moe_rejects_mismatched_expert_counts(): + with pytest.raises(ValueError, match="disagree on the number of experts"): + MoE(router=_make_router(), experts=_make_experts(num_experts=NUM_EXPERTS + 1)) + + +def test_moe_rejects_negative_aux_loss_coeff(): + with pytest.raises(ValueError, match="must be non-negative"): + MoE(router=_make_router(), experts=_make_experts(), aux_loss_coeff=-1.0) + + +def test_nemotron_3_nano_moe_dimensions(): + # Model report Table 1: 128 routed experts of dim 1856, 2 shared experts, top-6 routing. + # Megatron realizes the two shared experts as a single MLP of twice the expert dimension. + moe = MoE( + router=TopKRouter(n_embd=2688, num_experts=128, top_k=6, route_scale=2.5), + experts=GroupedExperts(n_embd=2688, ffn_hidden=1856, num_experts=128, backend=ExpertsBackend.LOOPED), + shared_experts=SquaredReLUMLP(n_embd=2688, ffn_hidden=2 * 1856), + aux_loss_coeff=1e-4, + ) + routed_params = sum(p.numel() for p in moe.experts.parameters()) + assert routed_params == 128 * 2 * 1856 * 2688 + # Active routed parameters per token: 6 of 128 experts. + assert routed_params * 6 // 128 == 6 * 2 * 1856 * 2688 + shared_params = sum(p.numel() for p in moe.shared_experts.parameters()) + assert shared_params == 2 * 3712 * 2688 diff --git a/tests/models/nemotron/test_moe_load_balancing.py b/tests/models/nemotron/test_moe_load_balancing.py new file mode 100644 index 000000000..7da73aed9 --- /dev/null +++ b/tests/models/nemotron/test_moe_load_balancing.py @@ -0,0 +1,340 @@ +"""Tests for auxiliary-loss-free MoE load balancing and the auxiliary-loss plumbing.""" + +import pytest +import torch +from torch.optim import AdamW + +from modalities.batch import InferenceResultBatch +from modalities.loss_functions import CLMCrossEntropyLoss +from modalities.models.components.moe.load_balancing import ( + MoEBalancing, + get_expert_load_reduction_group, + get_moe_layers, + update_expert_biases, +) +from modalities.models.components.moe.moe_losses import MoEAuxLoss, WeightedSumLoss +from modalities.models.nemotron.nemotron_model import NemotronLLM +from tests.models.nemotron.test_nemotron_model import VOCAB_SIZE, _layer_specs, _make_model + + +def _moe_model(aux_loss_coeff: float = 0.0, **overrides) -> NemotronLLM: + return _make_model( + layer_pattern="EE", + n_layer=2, + layer_specs=_layer_specs(aux_loss_coeff=aux_loss_coeff), + **overrides, + ) + + +# -------------------------------------------------------------------------------------------- +# Discovery +# -------------------------------------------------------------------------------------------- + + +def test_get_moe_layers_finds_layers_in_a_mixed_stack(): + model = _make_model(layer_pattern="MEE*-", n_layer=5) + assert len(get_moe_layers(model)) == 2 + + +def test_get_moe_layers_returns_empty_for_a_dense_model(): + assert get_moe_layers(_make_model(layer_pattern="M-", n_layer=2)) == [] + + +def test_reduction_group_is_none_without_distributed(): + # Single-process training needs no collective at all. + assert get_expert_load_reduction_group(device_mesh=None) is None + + +# -------------------------------------------------------------------------------------------- +# The bias update rule +# -------------------------------------------------------------------------------------------- + + +def test_update_moves_bias_up_for_underloaded_and_down_for_overloaded_experts(): + model = _moe_model() + moe = get_moe_layers(model)[0] + # Expert 0 is starved, expert 7 is swamped, the rest sit at the mean. + counts = torch.tensor([0.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 30.0]) + with torch.no_grad(): + moe.router.tokens_per_expert.copy_(counts) + + update_rate = 1e-3 + update_expert_biases([moe], update_rate=update_rate, process_group=None) + + bias = moe.router.expert_bias + assert bias[0].item() == pytest.approx(update_rate) + assert bias[7].item() == pytest.approx(-update_rate) + # Mean load is 11.25, so the experts at 10 are slightly under-loaded. + for expert_idx in range(1, 7): + assert bias[expert_idx].item() == pytest.approx(update_rate) + + +def test_update_uses_a_fixed_step_size_regardless_of_imbalance_magnitude(): + # The rule depends only on the sign of the deviation, which is what makes it robust to + # activation-checkpointing double counting and to gradient accumulation. + model = _moe_model() + moe = get_moe_layers(model)[0] + + for scale in (1.0, 1000.0): + with torch.no_grad(): + moe.router.expert_bias.zero_() + moe.router.tokens_per_expert.copy_(torch.tensor([0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]) * scale) + update_expert_biases([moe], update_rate=1e-3, process_group=None) + assert moe.router.expert_bias[0].item() == pytest.approx(1e-3) + + +def test_update_resets_the_token_counters(): + model = _moe_model() + moe = get_moe_layers(model)[0] + with torch.no_grad(): + moe.router.tokens_per_expert.fill_(5.0) + update_expert_biases([moe], update_rate=1e-3, process_group=None) + torch.testing.assert_close(moe.router.tokens_per_expert, torch.zeros_like(moe.router.tokens_per_expert)) + + +def test_update_is_a_noop_when_no_tokens_were_routed(): + model = _moe_model() + moe = get_moe_layers(model)[0] + with torch.no_grad(): + moe.router.expert_bias.fill_(0.5) + update_expert_biases([moe], update_rate=1e-3, process_group=None) + torch.testing.assert_close(moe.router.expert_bias, torch.full_like(moe.router.expert_bias, 0.5)) + + +def test_update_skips_layers_without_an_expert_bias(): + model = _make_model( + layer_pattern="E", + n_layer=1, + layer_specs={**_layer_specs(), "E": _layer_specs()["E"]}, + ) + moe = get_moe_layers(model)[0] + moe.router.expert_bias = None + with torch.no_grad(): + moe.router.tokens_per_expert.fill_(3.0) + # Must not raise, and must leave the counter alone (nothing to balance). + update_expert_biases([moe], update_rate=1e-3, process_group=None) + + +def test_repeated_updates_drive_the_load_towards_balance(): + # A small closed-loop simulation: the bias should progressively favour the starved expert. + model = _moe_model() + moe = get_moe_layers(model)[0] + starved = 3 + for _ in range(50): + counts = torch.full((8,), 10.0) + counts[starved] = 0.0 + with torch.no_grad(): + moe.router.tokens_per_expert.copy_(counts) + update_expert_biases([moe], update_rate=1e-2, process_group=None) + + bias = moe.router.expert_bias + assert bias[starved] > bias.mean() + assert bias[starved].item() == pytest.approx(50 * 1e-2, rel=1e-3) + + +def test_expert_bias_influences_routing_after_the_update(): + # The end-to-end effect: balancing must actually change which experts get picked. + model = _moe_model().eval() + moe = get_moe_layers(model)[0] + x = torch.randn(64, model.n_embd) + _, before, _ = moe.router(x) + + starved = int(torch.bincount(before.reshape(-1), minlength=8).argmin()) + with torch.no_grad(): + counts = torch.full((8,), 100.0) + counts[starved] = 0.0 + moe.router.tokens_per_expert.copy_(counts) + # A large step so that the effect is unambiguous. + update_expert_biases([moe], update_rate=10.0, process_group=None) + + _, after, _ = moe.router(x) + assert (after == starved).sum() > (before == starved).sum() + + +# -------------------------------------------------------------------------------------------- +# The optimizer decorator +# -------------------------------------------------------------------------------------------- + + +def test_decorator_returns_the_same_optimizer_instance(): + model = _moe_model() + optimizer = AdamW(model.parameters(), lr=1e-4) + decorated = MoEBalancing.register_expert_bias_update_hook( + optimizer=optimizer, model=model, expert_bias_update_rate=1e-3 + ) + assert decorated is optimizer + + +def test_hook_fires_once_per_optimizer_step_not_per_micro_batch(): + # Correctness under gradient accumulation: the bias must move by exactly one step size per + # optimizer step, no matter how many forward/backward passes happened in between. + model = _moe_model() + optimizer = AdamW(model.parameters(), lr=0.0) + MoEBalancing.register_expert_bias_update_hook(optimizer=optimizer, model=model, expert_bias_update_rate=1e-3) + moe = get_moe_layers(model)[0] + + gradient_accumulation_steps = 4 + for _ in range(gradient_accumulation_steps): + out = model({"input_ids": torch.randint(0, VOCAB_SIZE, (2, 8))}) + out["logits"].float().mean().backward() + # All four micro-batches have accumulated into the counter. + assert moe.router.tokens_per_expert.sum() == 2 * 8 * 2 * gradient_accumulation_steps + + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + # Exactly one update of magnitude update_rate was applied, and the counters were reset. + assert moe.router.expert_bias.abs().max().item() == pytest.approx(1e-3) + assert moe.router.tokens_per_expert.sum() == 0 + + +def test_hook_updates_every_moe_layer(): + model = _moe_model() + optimizer = AdamW(model.parameters(), lr=0.0) + MoEBalancing.register_expert_bias_update_hook(optimizer=optimizer, model=model, expert_bias_update_rate=1e-3) + model({"input_ids": torch.randint(0, VOCAB_SIZE, (2, 8))})["logits"].float().mean().backward() + optimizer.step() + + for moe in get_moe_layers(model): + assert moe.router.expert_bias.abs().max() > 0 + + +def test_decorator_warns_and_is_inert_without_moe_layers(caplog): + # A pipeline stage may legitimately hold no MoE layers. + model = _make_model(layer_pattern="M-", n_layer=2) + optimizer = AdamW(model.parameters(), lr=0.0) + with caplog.at_level("WARNING"): + MoEBalancing.register_expert_bias_update_hook(optimizer=optimizer, model=model, expert_bias_update_rate=1e-3) + assert "No MoE layers found" in caplog.text + optimizer.step() # must not raise + + +def test_decorator_warns_when_expert_bias_is_disabled(caplog): + specs = _layer_specs() + from modalities.models.components.moe.experts import ExpertsBackend + from modalities.models.nemotron.nemotron_layer_specs import NemotronMoELayerSpec + from tests.models.nemotron.test_nemotron_model import N_EMBD, NORM_CONFIG + + specs["E"] = NemotronMoELayerSpec( + n_embd=N_EMBD, + num_experts=8, + moe_ffn_hidden=32, + top_k=2, + use_expert_bias=False, + experts_backend=ExpertsBackend.LOOPED, + norm_config=NORM_CONFIG, + ) + model = _make_model(layer_pattern="E", n_layer=1, layer_specs=specs) + optimizer = AdamW(model.parameters(), lr=0.0) + with caplog.at_level("WARNING"): + MoEBalancing.register_expert_bias_update_hook(optimizer=optimizer, model=model, expert_bias_update_rate=1e-3) + assert "none of them maintains an expert bias" in caplog.text + + +def test_decorator_rejects_non_positive_update_rate(): + model = _moe_model() + optimizer = AdamW(model.parameters(), lr=1e-4) + with pytest.raises(ValueError, match="must be positive"): + MoEBalancing.register_expert_bias_update_hook(optimizer=optimizer, model=model, expert_bias_update_rate=0.0) + + +def test_expert_bias_survives_a_state_dict_round_trip(): + # The bias is training state that must be checkpointed, or balancing restarts from scratch on + # every warmstart. + model = _moe_model() + moe = get_moe_layers(model)[0] + with torch.no_grad(): + moe.router.expert_bias.copy_(torch.linspace(-1.0, 1.0, 8)) + + restored = _moe_model() + restored.load_state_dict(model.state_dict()) + torch.testing.assert_close(get_moe_layers(restored)[0].router.expert_bias, moe.router.expert_bias) + + +# -------------------------------------------------------------------------------------------- +# Auxiliary loss plumbing +# -------------------------------------------------------------------------------------------- + + +def _forward_batch(model: NemotronLLM, batch_size: int = 2, seq_len: int = 8) -> InferenceResultBatch: + inputs = torch.randint(0, VOCAB_SIZE, (batch_size, seq_len)) + targets = torch.randint(0, VOCAB_SIZE, (batch_size, seq_len)) + return InferenceResultBatch(targets={"target_ids": targets}, predictions=model({"input_ids": inputs})) + + +def test_moe_aux_loss_reads_the_value_from_the_forward_batch(): + model = _moe_model(aux_loss_coeff=1e-2, aux_loss_key="moe_aux_loss") + batch = _forward_batch(model) + aux_loss = MoEAuxLoss(prediction_key="moe_aux_loss") + torch.testing.assert_close(aux_loss(batch), model.get_aux_loss()) + + +def test_inference_result_batch_length_is_unaffected_by_the_scalar_aux_loss(): + # InferenceResultBatch derives its length from the *first* predictions entry, so the logits + # must be inserted before the scalar auxiliary loss. + model = _moe_model(aux_loss_coeff=1e-2, aux_loss_key="moe_aux_loss") + batch = _forward_batch(model, batch_size=3) + assert len(batch) == 3 + + +def test_weighted_sum_loss_combines_terms(): + model = _moe_model(aux_loss_coeff=1e-2, aux_loss_key="moe_aux_loss") + batch = _forward_batch(model) + + clm = CLMCrossEntropyLoss(target_key="target_ids", prediction_key="logits") + aux = MoEAuxLoss(prediction_key="moe_aux_loss") + combined = WeightedSumLoss(losses=[clm, aux], weights=[1.0, 1.0]) + + torch.testing.assert_close(combined(batch), clm(batch) + aux(batch), rtol=1e-5, atol=1e-7) + + +def test_weighted_sum_loss_applies_the_weights(): + model = _moe_model(aux_loss_coeff=1e-2, aux_loss_key="moe_aux_loss") + batch = _forward_batch(model) + clm = CLMCrossEntropyLoss(target_key="target_ids", prediction_key="logits") + aux = MoEAuxLoss(prediction_key="moe_aux_loss") + + weighted = WeightedSumLoss(losses=[clm, aux], weights=[0.5, 3.0]) + torch.testing.assert_close(weighted(batch), 0.5 * clm(batch) + 3.0 * aux(batch), rtol=1e-5, atol=1e-7) + + +def test_weighted_sum_loss_is_differentiable_into_the_router(): + model = _moe_model(aux_loss_coeff=1e-2, aux_loss_key="moe_aux_loss") + batch = _forward_batch(model) + combined = WeightedSumLoss( + losses=[ + CLMCrossEntropyLoss(target_key="target_ids", prediction_key="logits"), + MoEAuxLoss(prediction_key="moe_aux_loss"), + ], + weights=[1.0, 1.0], + ) + combined(batch).backward() + for moe in get_moe_layers(model): + assert moe.router.gate.weight.grad is not None + assert moe.router.gate.weight.grad.abs().sum() > 0 + + +def test_weighted_sum_loss_rejects_mismatched_weights(): + clm = CLMCrossEntropyLoss(target_key="target_ids", prediction_key="logits") + with pytest.raises(ValueError, match="they must match"): + WeightedSumLoss(losses=[clm], weights=[1.0, 2.0]) + with pytest.raises(ValueError, match="at least one loss"): + WeightedSumLoss(losses=[], weights=[]) + + +def test_moe_aux_loss_raises_a_clear_error_when_the_key_is_absent(): + # A likely misconfiguration: model.aux_loss_key not matching the loss's prediction_key. + model = _moe_model(aux_loss_coeff=1e-2, aux_loss_key="moe_aux_loss") + batch = _forward_batch(model) + with pytest.raises(Exception, match="not present in predictions"): + MoEAuxLoss(prediction_key="wrong_key")(batch) + + +def test_registry_exposes_the_phase_four_components(): + from modalities.registry.components import COMPONENTS + from modalities.registry.registry import Registry + + registry = Registry(COMPONENTS) + assert registry.get_component("loss", "moe_aux_loss") is MoEAuxLoss + assert registry.get_component("loss", "weighted_sum") is WeightedSumLoss + assert registry.get_component("optimizer", "moe_load_balanced") is not None diff --git a/tests/models/nemotron/test_nemotron_attention.py b/tests/models/nemotron/test_nemotron_attention.py new file mode 100644 index 000000000..0c4f9b3f1 --- /dev/null +++ b/tests/models/nemotron/test_nemotron_attention.py @@ -0,0 +1,166 @@ +import pytest +import torch + +from modalities.models.nemotron.nemotron_attention import NemotronAttentionImplementation, NemotronSelfAttention + +# head_dim is deliberately not n_embd // n_head_q: that decoupling is the reason this module +# exists instead of reusing the GPT2 attention. +ATTENTION_KWARGS = dict(n_embd=24, n_head_q=4, n_head_kv=2, head_dim=8) + + +def _make_attention(**overrides) -> NemotronSelfAttention: + torch.manual_seed(0) + return NemotronSelfAttention(**{**ATTENTION_KWARGS, **overrides}) + + +def test_projection_shapes_decouple_head_dim_from_model_dim(): + attn = _make_attention() + assert attn.q_attn.weight.shape == (4 * 8, 24) + assert attn.k_attn.weight.shape == (2 * 8, 24) + assert attn.v_attn.weight.shape == (2 * 8, 24) + assert attn.c_proj.weight.shape == (24, 4 * 8) + assert attn.q_attn.bias is None + assert attn.n_rep == 2 + + +def test_nemotron_3_nano_attention_dimensions(): + # Model report Table 1: 32 Q heads, 2 KV heads, head dim 128, model dim 2688. + attn = NemotronSelfAttention(n_embd=2688, n_head_q=32, n_head_kv=2, head_dim=128) + assert attn.q_attn.out_features == 4096 + assert attn.k_attn.out_features == 256 + assert attn.v_attn.out_features == 256 + assert attn.c_proj.in_features == 4096 + assert attn.n_rep == 16 + + +def test_forward_preserves_shape(): + attn = _make_attention() + x = torch.randn(2, 7, 24) + out = attn(x) + assert out.shape == x.shape + assert torch.isfinite(out).all() + + +def test_manual_and_pytorch_flash_implementations_agree(): + manual = _make_attention(attention_implementation=NemotronAttentionImplementation.MANUAL).eval() + flash = _make_attention(attention_implementation=NemotronAttentionImplementation.PYTORCH_FLASH).eval() + flash.load_state_dict(manual.state_dict()) + + x = torch.randn(2, 12, 24) + with torch.no_grad(): + torch.testing.assert_close(flash(x), manual(x), rtol=1e-4, atol=1e-5) + + +@pytest.mark.parametrize( + "implementation", + [NemotronAttentionImplementation.MANUAL, NemotronAttentionImplementation.PYTORCH_FLASH], +) +def test_attention_is_causal(implementation): + attn = _make_attention(attention_implementation=implementation).eval() + x = torch.randn(1, 10, 24) + with torch.no_grad(): + baseline = attn(x) + perturbed = x.clone() + perturbed[:, 6] += 5.0 + outputs = attn(perturbed) + + torch.testing.assert_close(outputs[:, :6], baseline[:, :6], rtol=1e-4, atol=1e-5) + assert not torch.allclose(outputs[:, 6], baseline[:, 6]) + + +def test_grouped_query_attention_matches_explicitly_repeated_kv_heads(): + # A GQA model with n_head_kv == n_head_q must behave exactly like multi-head attention; + # this pins the head-repetition logic against a configuration where it is a no-op. + gqa = _make_attention(n_head_q=4, n_head_kv=4).eval() + x = torch.randn(2, 6, 24) + with torch.no_grad(): + out = gqa(x) + assert gqa.n_rep == 1 + assert out.shape == x.shape + + +def test_repeat_kv_duplicates_each_head_contiguously(): + x = torch.arange(2 * 2 * 3 * 2, dtype=torch.float32).reshape(2, 2, 3, 2) + repeated = NemotronSelfAttention._repeat_kv(x, n_rep=3) + assert repeated.shape == (2, 6, 3, 2) + # Heads 0,1,2 are all copies of source head 0; heads 3,4,5 of source head 1. + for head in range(3): + torch.testing.assert_close(repeated[:, head], x[:, 0]) + for head in range(3, 6): + torch.testing.assert_close(repeated[:, head], x[:, 1]) + + +def test_repeat_kv_is_a_noop_for_n_rep_one(): + x = torch.randn(1, 2, 3, 4) + assert NemotronSelfAttention._repeat_kv(x, n_rep=1) is x + + +def test_batch_elements_are_independent(): + attn = _make_attention().eval() + x = torch.randn(3, 8, 24) + with torch.no_grad(): + batched = attn(x) + individually = torch.cat([attn(x[i : i + 1]) for i in range(3)], dim=0) + torch.testing.assert_close(batched, individually, rtol=1e-4, atol=1e-5) + + +def test_attention_is_differentiable(): + attn = _make_attention() + x = torch.randn(2, 6, 24, requires_grad=True) + attn(x).sum().backward() + assert x.grad is not None and torch.isfinite(x.grad).all() + for name, param in attn.named_parameters(): + assert param.grad is not None, f"{name} received no gradient" + + +def test_dropout_is_disabled_in_eval_mode(): + attn = _make_attention(dropout=0.9).eval() + x = torch.randn(2, 6, 24) + with torch.no_grad(): + torch.testing.assert_close(attn(x), attn(x)) + + +def test_attention_rejects_indivisible_head_counts(): + with pytest.raises(ValueError, match="must be divisible by n_head_kv"): + _make_attention(n_head_q=3, n_head_kv=2) + + +def test_dao_flash_raises_when_not_installed_or_matches_reference(): + try: + import flash_attn # noqa: F401 + except ModuleNotFoundError: + attn = _make_attention(attention_implementation=NemotronAttentionImplementation.DAO_FLASH) + with pytest.raises(NotImplementedError, match="flash-attn is not installed"): + attn(torch.randn(1, 4, 24)) + else: + if not torch.cuda.is_available(): + pytest.skip("flash-attn requires CUDA") + # head_dim must be a multiple of 8 for the Dao kernel; use a realistic configuration. + manual = ( + NemotronSelfAttention( + n_embd=64, + n_head_q=4, + n_head_kv=2, + head_dim=32, + attention_implementation=NemotronAttentionImplementation.MANUAL, + ) + .cuda() + .bfloat16() + .eval() + ) + dao = ( + NemotronSelfAttention( + n_embd=64, + n_head_q=4, + n_head_kv=2, + head_dim=32, + attention_implementation=NemotronAttentionImplementation.DAO_FLASH, + ) + .cuda() + .bfloat16() + .eval() + ) + dao.load_state_dict(manual.state_dict()) + x = torch.randn(2, 16, 64, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + torch.testing.assert_close(dao(x), manual(x), rtol=2e-2, atol=2e-2) diff --git a/tests/models/nemotron/test_nemotron_config_build.py b/tests/models/nemotron/test_nemotron_config_build.py new file mode 100644 index 000000000..ac75758d9 --- /dev/null +++ b/tests/models/nemotron/test_nemotron_config_build.py @@ -0,0 +1,147 @@ +"""End-to-end test of the YAML -> registry -> model path. + +This is the integration test that proves the layer-spec design actually works through Modalities' +component factory: nested ``nemotron_layer_spec`` components must be built as *builders*, and the +model must then turn them into one independent module per layer position. +""" + +from pathlib import Path + +import pytest +import torch +from pydantic import BaseModel + +from modalities.config.component_factory import ComponentFactory +from modalities.config.config import load_app_config_dict +from modalities.config.pydantic_if_types import PydanticPytorchModuleType +from modalities.models.nemotron.nemotron_layers import ( + Mamba2Layer, + NemotronAttentionLayer, + NemotronMLPLayer, + NemotronMoELayer, +) +from modalities.models.nemotron.nemotron_model import NemotronLLM +from modalities.registry.components import COMPONENTS +from modalities.registry.registry import Registry + +CONFIG_PATH = Path(__file__).parents[2] / "test_yaml_configs" / "nemotron_config_initialization.yaml" + + +class _RawModelConfig(BaseModel): + model_raw: PydanticPytorchModuleType + + +def _build_raw_model() -> NemotronLLM: + config_dict = load_app_config_dict(config_file_path=CONFIG_PATH) + # The initialization sub-config carries placeholders that only the FSDP tests replace; the raw + # model does not depend on them. + config_dict.pop("model", None) + component_factory = ComponentFactory(registry=Registry(COMPONENTS)) + components = component_factory.build_components(config_dict=config_dict, components_model_type=_RawModelConfig) + return components.model_raw + + +@pytest.fixture(scope="module") +def model() -> NemotronLLM: + return _build_raw_model() + + +def test_model_is_built_from_yaml(model): + assert isinstance(model, NemotronLLM) + assert model.n_layer == 8 + assert model.layer_pattern == "MEMEM*E-" + + +def test_layer_types_follow_the_configured_pattern(model): + expected = [ + Mamba2Layer, + NemotronMoELayer, + Mamba2Layer, + NemotronMoELayer, + Mamba2Layer, + NemotronAttentionLayer, + NemotronMoELayer, + NemotronMLPLayer, + ] + for layer_idx, expected_type in enumerate(expected): + assert isinstance(model.transformer.h[str(layer_idx)], expected_type), f"layer {layer_idx}" + + +def test_repeated_layer_types_have_independent_weights(model): + # The whole reason layer specs are builders: the component factory memoises components, so + # injecting instantiated layers would make all three Mamba layers share one weight tensor. + mamba_layers = [model.transformer.h[idx] for idx in ("0", "2", "4")] + weights = [layer.mixer.in_proj.weight for layer in mamba_layers] + assert len({id(w) for w in weights}) == 3 + assert not torch.equal(weights[0], weights[1]) + assert not torch.equal(weights[1], weights[2]) + + moe_layers = [model.transformer.h[idx] for idx in ("1", "3", "6")] + router_weights = [layer.moe.router.gate.weight for layer in moe_layers] + assert len({id(w) for w in router_weights}) == 3 + + +def test_component_hyperparameters_are_taken_from_yaml(model): + mixer = model.transformer.h["0"].mixer + assert (mixer.n_heads, mixer.head_dim, mixer.state_dim, mixer.n_groups) == (8, 32, 16, 2) + assert mixer.d_inner == 8 * 32 + assert mixer.chunk_size == 32 + + moe = model.transformer.h["1"].moe + assert moe.router.num_experts == 8 + assert moe.router.top_k == 2 + assert moe.router.route_scale == 2.5 + assert moe.router.expert_bias is not None + assert moe.experts.ffn_hidden == 64 + # 2 shared experts fused into one MLP of 2 * moe_ffn_hidden hidden units. + assert moe.shared_experts.c_fc.out_features == 128 + assert moe.aux_loss_coeff == pytest.approx(1e-4) + + attn = model.transformer.h["5"].attn + assert (attn.n_head_q, attn.n_head_kv, attn.head_dim) == (8, 2, 32) + + mlp = model.transformer.h["7"].mlp + assert mlp.c_fc.out_features == 128 + + +def test_forward_pass_through_the_configured_model(model): + inputs = {"input_ids": torch.randint(0, model.vocab_size, (2, 16))} + out = model(inputs) + assert list(out.keys()) == ["logits", "moe_aux_loss"] + assert out["logits"].shape == (2, 16, model.vocab_size) + assert torch.isfinite(out["logits"]).all() + assert torch.isfinite(out["moe_aux_loss"]) + + +def test_backward_pass_through_the_configured_model(model): + out = model({"input_ids": torch.randint(0, model.vocab_size, (2, 16))}) + (out["logits"].float().mean() + out["moe_aux_loss"]).backward() + assert model.transformer.wte.weight.grad is not None + assert model.transformer.h["1"].moe.router.gate.weight.grad is not None + model.zero_grad(set_to_none=True) + + +def test_activation_checkpointing_applies_to_the_layer_stack(model): + # The Modalities activation checkpointing component requires transformer.h to be a ModuleDict + # and addresses it by fully qualified name. + from modalities.config.config import ActivationCheckpointedModelConfig + from modalities.models.model_factory import ModelFactory + from modalities.training.activation_checkpointing.activation_checkpointing_variants import ( + ActivationCheckpointingVariants, + ) + + checkpointed = ModelFactory.get_activation_checkpointed_fsdp2_model_( + ac_variant=ActivationCheckpointingVariants.FULL_ACTIVATION_CHECKPOINTING, + layers_fqn="transformer.h", + model=_build_raw_model(), + ac_fun_params=ActivationCheckpointedModelConfig.FullACParams(), + ) + out = checkpointed({"input_ids": torch.randint(0, model.vocab_size, (1, 8))}) + assert out["logits"].shape == (1, 8, model.vocab_size) + + +def test_fsdp2_block_names_resolve_to_the_layer_classes(model): + from modalities.util import get_module_class_from_name + + for block_name in ("Mamba2Layer", "NemotronMoELayer", "NemotronAttentionLayer", "NemotronMLPLayer"): + assert get_module_class_from_name(model, block_name) is not None, block_name diff --git a/tests/models/nemotron/test_nemotron_initialization.py b/tests/models/nemotron/test_nemotron_initialization.py new file mode 100644 index 000000000..9b1f9baa0 --- /dev/null +++ b/tests/models/nemotron/test_nemotron_initialization.py @@ -0,0 +1,247 @@ +"""Weight-initialization tests for the Nemotron hybrid model. + +The interesting property here is *separation of concerns*: the generic, regex-driven initializer +handles all linear and embedding weights, while the state space parameters (A_log, D, dt_bias, +conv1d) follow their own distributions defined in ``Mamba2Mixer.reset_parameters``. If a regex ever +matched an SSM parameter, the model would still train but with a silently broken state space. +""" + +import re + +import pytest +import torch + +from modalities.models.components.moe.experts import ExpertsBackend, GroupedExperts +from modalities.models.model_factory import ModelFactory +from modalities.nn.model_initialization.composed_initialization import ComposedInitializationRoutines +from modalities.nn.model_initialization.parameter_name_filters import ( + NAMED_PARAMETER_INIT_GROUPS, + SupportWeightInitModels, + WeightInitTypes, +) +from tests.models.nemotron.test_nemotron_model import _make_model + +SSM_PARAMETER_SUFFIXES = ("A_log", "D", "dt_bias", "conv1d_weight", "conv1d_bias") + + +def _all_init_regexes() -> list[str]: + groups = NAMED_PARAMETER_INIT_GROUPS[SupportWeightInitModels.NEMOTRON] + regexes: list[str] = [] + for weight_init_type in WeightInitTypes: + regex_filter = groups[weight_init_type] + regexes.extend(regex_filter.weights) + regexes.extend(regex_filter.biases or []) + return regexes + + +def test_nemotron_is_a_supported_weight_init_model(): + assert SupportWeightInitModels("nemotron") is SupportWeightInitModels.NEMOTRON + groups = NAMED_PARAMETER_INIT_GROUPS[SupportWeightInitModels.NEMOTRON] + for weight_init_type in WeightInitTypes: + assert groups[weight_init_type] is not None + + +def test_plain_init_covers_every_linear_and_embedding_weight(): + model = _make_model() + plain = NAMED_PARAMETER_INIT_GROUPS[SupportWeightInitModels.NEMOTRON][WeightInitTypes.PLAIN] + patterns = plain.weights + + for name, param in model.named_parameters(): + is_ssm = name.endswith(SSM_PARAMETER_SUFFIXES) + is_norm = ".norm." in name or "lm_head_norm" in name + if is_ssm or is_norm: + continue + assert any(re.fullmatch(p, name) for p in patterns), f"{name} is not covered by the plain init filter" + + +def test_no_init_regex_matches_a_state_space_parameter(): + # The load-bearing assertion of this file. + model = _make_model() + regexes = _all_init_regexes() + for name, _ in model.named_parameters(): + if name.endswith(SSM_PARAMETER_SUFFIXES): + matching = [p for p in regexes if re.fullmatch(p, name)] + assert not matching, f"SSM parameter {name} would be overwritten by {matching}" + + +def test_no_init_regex_matches_a_normalization_weight(): + # Norms are initialized to one at construction time and must stay that way. + model = _make_model() + regexes = _all_init_regexes() + for name, _ in model.named_parameters(): + if ".norm." in name or "lm_head_norm" in name: + matching = [p for p in regexes if re.fullmatch(p, name)] + assert not matching, f"Normalization parameter {name} would be overwritten by {matching}" + + +def test_scaled_init_covers_every_residual_output_projection(): + scaled = NAMED_PARAMETER_INIT_GROUPS[SupportWeightInitModels.NEMOTRON][WeightInitTypes.SCALED] + model = _make_model() + matched = {name for name, _ in model.named_parameters() if any(re.fullmatch(p, name) for p in scaled.weights)} + # One per layer: the projection that writes back into the residual stream. + assert matched == { + "transformer.h.0.mixer.out_proj.weight", + "transformer.h.1.moe.experts.w2", + "transformer.h.1.moe.shared_experts.c_proj.weight", + "transformer.h.2.attn.c_proj.weight", + "transformer.h.3.mlp.c_proj.weight", + } + + +def test_composed_initializer_runs_and_preserves_ssm_distributions(): + model = _make_model() + ssm_before = { + name: param.detach().clone() + for name, param in model.named_parameters() + if name.endswith(SSM_PARAMETER_SUFFIXES) + } + assert ssm_before, "test model must contain a Mamba layer" + + initializer = ComposedInitializationRoutines.get_composed_model_initializer( + model_type=SupportWeightInitModels.NEMOTRON, + weight_init_type=WeightInitTypes.SCALED, + mean=0.0, + std=0.02, + num_layers=model.n_layer, + seed=42, + ) + initializer.initialize_in_place(model) + + for name, before in ssm_before.items(): + torch.testing.assert_close(dict(model.named_parameters())[name], before, msg=f"{name} was modified") + + +def test_composed_initializer_sets_linear_weights_to_the_requested_std(): + model = _make_model() + std = 0.02 + initializer = ComposedInitializationRoutines.get_composed_model_initializer( + model_type=SupportWeightInitModels.NEMOTRON, + weight_init_type=WeightInitTypes.PLAIN, + mean=0.0, + std=std, + seed=42, + ) + initializer.initialize_in_place(model) + + # The embedding is the largest single tensor, so its empirical std is the most reliable. + embedding_std = model.transformer.wte.weight.std().item() + assert embedding_std == pytest.approx(std, rel=0.1) + + +def test_scaled_init_shrinks_residual_projections_relative_to_plain(): + std, num_layers = 0.02, 4 + plain_model = _make_model() + scaled_model = _make_model() + + for weight_init_type, model in ((WeightInitTypes.PLAIN, plain_model), (WeightInitTypes.SCALED, scaled_model)): + initializer = ComposedInitializationRoutines.get_composed_model_initializer( + model_type=SupportWeightInitModels.NEMOTRON, + weight_init_type=weight_init_type, + mean=0.0, + std=std, + num_layers=num_layers if weight_init_type == WeightInitTypes.SCALED else None, + seed=42, + ) + initializer.initialize_in_place(model) + + plain_proj = plain_model.transformer.h["0"].mixer.out_proj.weight.std().item() + scaled_proj = scaled_model.transformer.h["0"].mixer.out_proj.weight.std().item() + # Scaled init divides by sqrt(2 * num_layers). + assert scaled_proj == pytest.approx(plain_proj / (2 * num_layers) ** 0.5, rel=0.15) + + +def test_grouped_experts_are_initialized_at_construction(): + # Regression guard: w1/w2 are raw nn.Parameter(torch.empty(...)) tensors, so without an + # explicit reset_parameters they would contain uninitialized memory. + experts = GroupedExperts(n_embd=16, ffn_hidden=24, num_experts=4, backend=ExpertsBackend.LOOPED) + for weight in (experts.w1, experts.w2): + assert torch.isfinite(weight).all() + assert weight.abs().max() < 1.0 + + +def test_grouped_experts_reset_parameters_is_a_noop_on_meta_device(): + with torch.device("meta"): + experts = GroupedExperts(n_embd=16, ffn_hidden=24, num_experts=4, backend=ExpertsBackend.LOOPED) + experts.reset_parameters() + assert experts.w1.is_meta + + +def test_freshly_built_model_has_no_uninitialized_parameters(): + # Every parameter must be finite straight after construction, before any initializer runs. + model = _make_model() + for name, param in model.named_parameters(): + assert torch.isfinite(param).all(), f"{name} contains non-finite values after construction" + for name, buffer in model.named_buffers(): + assert torch.isfinite(buffer).all(), f"buffer {name} contains non-finite values after construction" + + +def test_meta_device_model_is_fully_initialized_by_the_model_factory(): + # This is the real production path: build on meta, then let ModelFactory materialize, run + # reset_parameters recursively and apply the configured initializer. + from modalities.models.nemotron.nemotron_model_factory import NemotronModelFactory + from tests.models.nemotron.test_nemotron_model import _model_kwargs + + model = NemotronModelFactory.get_nemotron_model(**_model_kwargs(), use_meta_device=True) + assert all(p.is_meta for p in model.parameters()) + + initializer = ComposedInitializationRoutines.get_composed_model_initializer( + model_type=SupportWeightInitModels.NEMOTRON, + weight_init_type=WeightInitTypes.PLAIN, + mean=0.0, + std=0.02, + seed=42, + ) + if not torch.cuda.is_available(): + pytest.skip("ModelFactory.get_weight_initialized_model materializes onto CUDA") + + model = ModelFactory.get_weight_initialized_model(model=model, model_initializer=initializer) + for name, param in model.named_parameters(): + assert torch.isfinite(param).all(), f"{name} is not finite after initialization" + + # The SSM parameters must have their own distributions, not the normal one. + mixer = model.transformer.h["0"].mixer + A = torch.exp(mixer.A_log) + assert A.min() >= 1.0 and A.max() <= 16.0 + torch.testing.assert_close(mixer.D, torch.ones_like(mixer.D)) + + +def test_initializer_matches_parameters_behind_wrapper_modules(): + # Regression guard for a silent-correctness bug: activation checkpointing (and torch.compile, + # and FSDP1) insert wrapper segments into parameter FQNs. The initialization filters are written + # against the plain model, so those segments must be stripped before matching - otherwise every + # per-layer regex silently fails and the model keeps its default initialization. + from modalities.config.config import ActivationCheckpointedModelConfig + from modalities.models.model_factory import ModelFactory + from modalities.nn.model_initialization.initialization_routines import normalize_parameter_name + from modalities.training.activation_checkpointing.activation_checkpointing_variants import ( + ActivationCheckpointingVariants, + ) + + assert ( + normalize_parameter_name("transformer.h.0._checkpoint_wrapped_module.mixer.in_proj.weight") + == "transformer.h.0.mixer.in_proj.weight" + ) + assert normalize_parameter_name("_orig_mod.transformer.wte.weight") == "transformer.wte.weight" + + model = ModelFactory.get_activation_checkpointed_fsdp2_model_( + ac_variant=ActivationCheckpointingVariants.FULL_ACTIVATION_CHECKPOINTING, + layers_fqn="transformer.h", + model=_make_model(), + ac_fun_params=ActivationCheckpointedModelConfig.FullACParams(), + ) + # The wrapper really is in the names, so this test would pass trivially otherwise. + assert any("_checkpoint_wrapped_module" in name for name, _ in model.named_parameters()) + + std = 0.02 + initializer = ComposedInitializationRoutines.get_composed_model_initializer( + model_type=SupportWeightInitModels.NEMOTRON, + weight_init_type=WeightInitTypes.PLAIN, + mean=0.0, + std=std, + seed=42, + ) + initializer.initialize_in_place(model) + + # A per-layer parameter behind the wrapper must have been initialized to the requested std. + in_proj = model.get_submodule("transformer.h.0").mixer.in_proj.weight + assert in_proj.std().item() == pytest.approx(std, rel=0.15) diff --git a/tests/models/nemotron/test_nemotron_layers.py b/tests/models/nemotron/test_nemotron_layers.py new file mode 100644 index 000000000..83a4103d3 --- /dev/null +++ b/tests/models/nemotron/test_nemotron_layers.py @@ -0,0 +1,141 @@ +import pytest +import torch +import torch.nn as nn + +from modalities.models.components.mamba2.mamba2_mixer import Mamba2Mixer +from modalities.models.components.moe.experts import ExpertsBackend, GroupedExperts +from modalities.models.components.moe.moe import MoE +from modalities.models.components.moe.router import TopKRouter +from modalities.models.nemotron.nemotron_attention import NemotronSelfAttention +from modalities.models.nemotron.nemotron_layers import ( + Mamba2Layer, + NemotronAttentionLayer, + NemotronMLPLayer, + NemotronMoELayer, +) +from modalities.models.nemotron.nemotron_mlp import SquaredReLUMLP + +N_EMBD = 16 + + +def _norm() -> nn.Module: + return nn.RMSNorm(normalized_shape=N_EMBD, eps=1e-5) + + +def _mamba_layer() -> Mamba2Layer: + torch.manual_seed(0) + return Mamba2Layer( + norm=_norm(), + mixer=Mamba2Mixer(n_embd=N_EMBD, n_heads=4, head_dim=8, state_dim=4, n_groups=2, chunk_size=4), + ) + + +def _attention_layer() -> NemotronAttentionLayer: + torch.manual_seed(0) + return NemotronAttentionLayer( + norm=_norm(), + attn=NemotronSelfAttention(n_embd=N_EMBD, n_head_q=4, n_head_kv=2, head_dim=8), + ) + + +def _moe_layer() -> NemotronMoELayer: + torch.manual_seed(0) + experts = GroupedExperts(n_embd=N_EMBD, ffn_hidden=24, num_experts=4, backend=ExpertsBackend.LOOPED) + with torch.no_grad(): + experts.w1.normal_(0, 0.02) + experts.w2.normal_(0, 0.02) + return NemotronMoELayer( + norm=_norm(), + moe=MoE( + router=TopKRouter(n_embd=N_EMBD, num_experts=4, top_k=2), + experts=experts, + shared_experts=SquaredReLUMLP(n_embd=N_EMBD, ffn_hidden=48), + ), + ) + + +def _mlp_layer() -> NemotronMLPLayer: + torch.manual_seed(0) + return NemotronMLPLayer(norm=_norm(), mlp=SquaredReLUMLP(n_embd=N_EMBD, ffn_hidden=24)) + + +ALL_LAYER_BUILDERS = [_mamba_layer, _attention_layer, _moe_layer, _mlp_layer] + + +@pytest.mark.parametrize("build_layer", ALL_LAYER_BUILDERS, ids=lambda f: f.__name__) +def test_layer_preserves_shape(build_layer): + layer = build_layer() + x = torch.randn(2, 8, N_EMBD) + out = layer(x) + assert out.shape == x.shape + assert torch.isfinite(out).all() + + +@pytest.mark.parametrize("build_layer", ALL_LAYER_BUILDERS, ids=lambda f: f.__name__) +def test_layer_is_a_residual_around_its_operator(build_layer): + # Every layer must be exactly x + operator(norm(x)); a missing residual is a silent + # convergence bug rather than a crash, so assert the identity directly. + layer = build_layer().eval() + x = torch.randn(2, 6, N_EMBD) + with torch.no_grad(): + expected = x + layer._operator(layer.norm(x)) + torch.testing.assert_close(layer(x), expected, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize("build_layer", ALL_LAYER_BUILDERS, ids=lambda f: f.__name__) +def test_layer_owns_exactly_one_norm(build_layer): + layer = build_layer() + norms = [module for module in layer.modules() if isinstance(module, (nn.RMSNorm, nn.LayerNorm))] + # The Mamba mixer contains its own internal gated norm, which is not an nn.RMSNorm. + assert len(norms) == 1 + assert norms[0] is layer.norm + + +@pytest.mark.parametrize("build_layer", ALL_LAYER_BUILDERS, ids=lambda f: f.__name__) +def test_layer_is_differentiable(build_layer): + layer = build_layer() + x = torch.randn(2, 6, N_EMBD, requires_grad=True) + layer(x).sum().backward() + assert x.grad is not None and torch.isfinite(x.grad).all() + assert layer.norm.weight.grad is not None + + +@pytest.mark.parametrize("build_layer", [_mamba_layer, _attention_layer], ids=lambda f: f.__name__) +def test_sequence_mixing_layers_are_causal(build_layer): + layer = build_layer().eval() + x = torch.randn(1, 12, N_EMBD) + with torch.no_grad(): + baseline = layer(x) + perturbed = x.clone() + perturbed[:, 7] += 5.0 + outputs = layer(perturbed) + torch.testing.assert_close(outputs[:, :7], baseline[:, :7], rtol=1e-4, atol=1e-5) + + +@pytest.mark.parametrize("build_layer", [_moe_layer, _mlp_layer], ids=lambda f: f.__name__) +def test_feed_forward_layers_act_position_wise(build_layer): + # MoE and MLP layers must not mix across the sequence at all: changing one position may only + # change that position's output. + layer = build_layer().eval() + x = torch.randn(1, 10, N_EMBD) + with torch.no_grad(): + baseline = layer(x) + perturbed = x.clone() + perturbed[:, 4] += 3.0 + outputs = layer(perturbed) + + untouched = [i for i in range(10) if i != 4] + torch.testing.assert_close(outputs[:, untouched], baseline[:, untouched], rtol=1e-4, atol=1e-5) + assert not torch.allclose(outputs[:, 4], baseline[:, 4]) + + +def test_layer_names_are_stable_for_parameter_name_filters(): + # The weight-initialization and weight-decay filters match on parameter names, so the + # submodule attribute names are part of the public contract. + assert "mixer.in_proj.weight" in dict(_mamba_layer().named_parameters()) + assert "attn.q_attn.weight" in dict(_attention_layer().named_parameters()) + moe_params = dict(_moe_layer().named_parameters()) + assert "moe.router.gate.weight" in moe_params + assert "moe.experts.w1" in moe_params + assert "moe.shared_experts.c_fc.weight" in moe_params + assert "mlp.c_fc.weight" in dict(_mlp_layer().named_parameters()) diff --git a/tests/models/nemotron/test_nemotron_model.py b/tests/models/nemotron/test_nemotron_model.py new file mode 100644 index 000000000..0e924abe0 --- /dev/null +++ b/tests/models/nemotron/test_nemotron_model.py @@ -0,0 +1,522 @@ +import pytest +import torch +import torch.nn as nn + +from modalities.models.components.mamba2.mamba2_mixer import Mamba2Mixer +from modalities.models.components.moe.experts import ExpertsBackend +from modalities.models.components.moe.moe import MoE +from modalities.models.components.norms import NormWrapperConfig +from modalities.models.nemotron.layer_pattern import LayerSymbol +from modalities.models.nemotron.nemotron_attention import NemotronSelfAttention +from modalities.models.nemotron.nemotron_layer_specs import ( + Mamba2LayerSpec, + NemotronAttentionLayerSpec, + NemotronMLPLayerSpec, + NemotronMoELayerSpec, +) +from modalities.models.nemotron.nemotron_layers import ( + Mamba2Layer, + NemotronAttentionLayer, + NemotronMLPLayer, + NemotronMoELayer, +) +from modalities.models.nemotron.nemotron_model import NemotronLLM, NemotronLLMConfig +from modalities.models.nemotron.nemotron_model_factory import NemotronModelFactory + +N_EMBD = 128 +VOCAB_SIZE = 256 +SEQUENCE_LENGTH = 32 +# Exercises all four layer types. +LAYER_PATTERN = "ME*-" + +NORM_CONFIG = {"norm_type": "pytorch_rms_norm", "config": {"normalized_shape": N_EMBD, "eps": 1e-5}} + + +def _layer_specs(aux_loss_coeff: float = 0.0, num_shared_experts: int = 2) -> dict: + return { + "M": Mamba2LayerSpec( + n_embd=N_EMBD, + mamba_n_heads=8, + mamba_head_dim=16, + mamba_state_dim=8, + mamba_n_groups=2, + chunk_size=8, + norm_config=NORM_CONFIG, + ), + "E": NemotronMoELayerSpec( + n_embd=N_EMBD, + num_experts=8, + moe_ffn_hidden=32, + top_k=2, + route_scale=2.5, + num_shared_experts=num_shared_experts, + aux_loss_coeff=aux_loss_coeff, + experts_backend=ExpertsBackend.LOOPED, + norm_config=NORM_CONFIG, + ), + "*": NemotronAttentionLayerSpec( + n_embd=N_EMBD, + n_head_q=4, + n_head_kv=2, + head_dim=16, + attention_implementation="manual", + norm_config=NORM_CONFIG, + ), + "-": NemotronMLPLayerSpec(n_embd=N_EMBD, ffn_hidden=32, norm_config=NORM_CONFIG), + } + + +def _model_kwargs(**overrides) -> dict: + kwargs = dict( + sample_key="input_ids", + prediction_key="logits", + sequence_length=SEQUENCE_LENGTH, + vocab_size=VOCAB_SIZE, + n_embd=N_EMBD, + n_layer=len(LAYER_PATTERN), + layer_pattern=LAYER_PATTERN, + layer_specs=_layer_specs(), + lm_head_norm_config=NORM_CONFIG, + ) + kwargs.update(overrides) + return kwargs + + +def _make_model(**overrides) -> NemotronLLM: + torch.manual_seed(0) + kwargs = _model_kwargs(**overrides) + kwargs["lm_head_norm_config"] = NormWrapperConfig.model_validate(kwargs["lm_head_norm_config"]) + return NemotronLLM(**kwargs) + + +# -------------------------------------------------------------------------------------------- +# Structure +# -------------------------------------------------------------------------------------------- + + +def test_module_tree_mirrors_gpt2_layout(): + # The wrapper components (activation checkpointing, FSDP block names, pipeline splitting) + # address the model by these names, so the layout is part of the contract. + model = _make_model() + assert isinstance(model.transformer, nn.ModuleDict) + assert isinstance(model.transformer.wte, nn.Embedding) + # Activation checkpointing requires a ModuleDict at transformer.h. + assert isinstance(model.transformer.h, nn.ModuleDict) + assert isinstance(model.transformer.lm_head, nn.Linear) + assert model.get_submodule("transformer.h") is model.transformer.h + + +def test_layers_are_built_according_to_the_pattern(): + model = _make_model() + expected_types = [Mamba2Layer, NemotronMoELayer, NemotronAttentionLayer, NemotronMLPLayer] + for layer_idx, expected_type in enumerate(expected_types): + assert isinstance(model.transformer.h[str(layer_idx)], expected_type) + + +def test_repeated_layer_types_do_not_share_parameters(): + # This is the central reason layer specs are builders rather than instances: two Mamba layers + # must have independent weights. + model = _make_model(layer_pattern="MM", n_layer=2) + first, second = model.transformer.h["0"], model.transformer.h["1"] + assert first is not second + assert first.mixer.in_proj.weight is not second.mixer.in_proj.weight + assert not torch.equal(first.mixer.in_proj.weight, second.mixer.in_proj.weight) + assert first.norm.weight is not second.norm.weight + + +def test_each_layer_owns_its_own_norm(): + model = _make_model() + norm_ids = {id(model.transformer.h[idx].norm) for idx in model.transformer.h} + assert len(norm_ids) == len(LAYER_PATTERN) + + +def test_embeddings_are_untied_by_default(): + model = _make_model() + assert not model.has_tied_word_embeddings + assert model.transformer.wte.weight is not model.transformer.lm_head.weight + + +def test_weight_tying_can_be_enabled(): + model = _make_model(use_weight_tying=True) + assert model.has_tied_word_embeddings + assert model.transformer.wte.weight is model.transformer.lm_head.weight + + +def test_lm_head_property_and_skip_toggle(): + model = _make_model() + assert model.lm_head is model.transformer.lm_head + + inputs = torch.randint(0, VOCAB_SIZE, (2, 8)) + logits = model.forward_impl(inputs) + assert logits.shape == (2, 8, VOCAB_SIZE) + + model.set_skip_lm_head(True) + hidden = model.forward_impl(inputs) + assert hidden.shape == (2, 8, N_EMBD) + # The chunked loss applies the head itself, so the two must compose back to the logits. + torch.testing.assert_close(model.lm_head(hidden), logits, rtol=1e-4, atol=1e-5) + + +# -------------------------------------------------------------------------------------------- +# Forward pass +# -------------------------------------------------------------------------------------------- + + +def test_forward_with_dict_input_returns_logits_under_prediction_key(): + model = _make_model() + inputs = {"input_ids": torch.randint(0, VOCAB_SIZE, (2, 8))} + out = model(inputs) + assert list(out.keys()) == ["logits"] + assert out["logits"].shape == (2, 8, VOCAB_SIZE) + assert torch.isfinite(out["logits"]).all() + + +def test_forward_with_tensor_input_returns_logits_tensor(): + # Pipeline parallelism passes raw tensors between stages. + model = _make_model() + out = model(torch.randint(0, VOCAB_SIZE, (2, 8))) + assert isinstance(out, torch.Tensor) + assert out.shape == (2, 8, VOCAB_SIZE) + + +def test_forward_exposes_aux_loss_when_configured(): + model = _make_model(layer_specs=_layer_specs(aux_loss_coeff=1e-2), aux_loss_key="moe_aux_loss") + out = model({"input_ids": torch.randint(0, VOCAB_SIZE, (2, 8))}) + # Logits must stay first: InferenceResultBatch derives its batch length from the first entry. + assert list(out.keys()) == ["logits", "moe_aux_loss"] + assert out["moe_aux_loss"].ndim == 0 + + +def test_forward_omits_aux_loss_when_coefficient_is_zero(): + model = _make_model(aux_loss_key="moe_aux_loss") + out = model({"input_ids": torch.randint(0, VOCAB_SIZE, (2, 8))}) + assert "moe_aux_loss" not in out + + +def test_get_aux_loss_sums_over_moe_layers(): + model = _make_model(layer_pattern="EE", n_layer=2, layer_specs=_layer_specs(aux_loss_coeff=1e-2)) + model({"input_ids": torch.randint(0, VOCAB_SIZE, (2, 8))}) + per_layer = [model.transformer.h[idx].moe.last_aux_loss for idx in model.transformer.h] + torch.testing.assert_close(model.get_aux_loss(), torch.stack(per_layer).sum()) + + +def test_get_aux_loss_is_none_without_moe_layers(): + model = _make_model(layer_pattern="M-", n_layer=2) + model({"input_ids": torch.randint(0, VOCAB_SIZE, (1, 4))}) + assert model.get_aux_loss() is None + + +def test_get_moe_layers_finds_every_moe_block(): + model = _make_model(layer_pattern="MEE*", n_layer=4) + moe_layers = model.get_moe_layers() + assert len(moe_layers) == 2 + assert all(isinstance(layer, MoE) for layer in moe_layers) + + +def test_model_is_causal_end_to_end(): + # A causality bug anywhere in the stack (conv, scan, attention mask) leaks the target and + # produces a suspiciously low loss, so assert it on the assembled model too. + model = _make_model().eval() + inputs = torch.randint(0, VOCAB_SIZE, (1, 16)) + with torch.no_grad(): + baseline = model.forward_impl(inputs) + perturbed = inputs.clone() + perturbed[:, 11] = (perturbed[:, 11] + 7) % VOCAB_SIZE + outputs = model.forward_impl(perturbed) + torch.testing.assert_close(outputs[:, :11], baseline[:, :11], rtol=1e-4, atol=1e-4) + assert not torch.allclose(outputs[:, 11], baseline[:, 11]) + + +def test_forward_rejects_sequences_longer_than_configured(): + model = _make_model() + with pytest.raises(ValueError, match="maximum input sequence length"): + model.forward_impl(torch.randint(0, VOCAB_SIZE, (1, SEQUENCE_LENGTH + 1))) + + +def test_model_is_differentiable_everywhere_except_unrouted_experts(): + model = _make_model() + out = model({"input_ids": torch.randint(0, VOCAB_SIZE, (2, 8))}) + out["logits"].sum().backward() + + # Expert weights are sparse by construction; every other parameter must receive a gradient. + sparse_params = {"transformer.h.1.moe.experts.w1", "transformer.h.1.moe.experts.w2"} + for name, param in model.named_parameters(): + assert param.grad is not None, f"{name} received no gradient" + assert torch.isfinite(param.grad).all(), f"{name} has non-finite gradient" + if name not in sparse_params: + assert param.grad.abs().sum() > 0, f"{name} has an all-zero gradient" + + +# -------------------------------------------------------------------------------------------- +# Weight decay groups +# -------------------------------------------------------------------------------------------- + + +def test_weight_decay_groups_partition_all_parameters(): + import re + + model = _make_model() + groups = model.weight_decay_groups + assigned: dict[str, list[str]] = {} + for name, _ in model.named_parameters(): + matches = [group for group, patterns in groups.items() if any(re.search(p, name) for p in patterns)] + assert len(matches) == 1, f"{name} matched {matches}, expected exactly one weight decay group" + assigned.setdefault(matches[0], []).append(name) + + # Every declared group must be non-empty for this pattern, otherwise the optimizer factory + # would raise when the group is excluded from weight decay. + for group in groups: + assert assigned.get(group), f"weight decay group '{group}' is empty" + + +def test_ssm_parameters_are_in_their_own_weight_decay_group(): + import re + + model = _make_model() + ssm_patterns = model.weight_decay_groups["ssm"] + ssm_params = [name for name, _ in model.named_parameters() if any(re.search(p, name) for p in ssm_patterns)] + # A_log, D, dt_bias, conv1d_weight, conv1d_bias of the single Mamba layer. + assert sorted(ssm_params) == [ + "transformer.h.0.mixer.A_log", + "transformer.h.0.mixer.D", + "transformer.h.0.mixer.conv1d_bias", + "transformer.h.0.mixer.conv1d_weight", + "transformer.h.0.mixer.dt_bias", + ] + + +# -------------------------------------------------------------------------------------------- +# Parameter counts +# -------------------------------------------------------------------------------------------- + + +def test_parameter_count_matches_the_analytic_formula(): + model = _make_model(layer_pattern="ME*-", n_layer=4) + actual = sum(p.numel() for p in model.parameters()) + + embedding = VOCAB_SIZE * N_EMBD + lm_head = N_EMBD * VOCAB_SIZE + lm_head_norm = N_EMBD + per_layer_norm = N_EMBD + + d_inner = 8 * 16 + group_state = 2 * 8 + mamba = ( + N_EMBD * (2 * d_inner + 2 * group_state + 8) # in_proj + + (d_inner + 2 * group_state) * 4 # conv1d weight + + (d_inner + 2 * group_state) # conv1d bias + + 8 * 3 # A_log, D, dt_bias + + d_inner # mixer gated norm + + d_inner * N_EMBD # out_proj + ) + moe = ( + N_EMBD * 8 # router gate + + 8 * 2 * 32 * N_EMBD # 8 routed experts, 2 matrices each + + 2 * (2 * 32) * N_EMBD # fused shared expert MLP of hidden 2 * 32 + ) + attention = (4 * 16 + 2 * 16 + 2 * 16) * N_EMBD + 4 * 16 * N_EMBD + mlp = 2 * 32 * N_EMBD + + expected = embedding + lm_head + lm_head_norm + 4 * per_layer_norm + mamba + moe + attention + mlp + assert actual == expected + + +def test_nemotron_3_nano_parameter_count_matches_the_model_report(): + # Model report: 31.6B total parameters, 3.2B active (3.6B including embeddings). + # Reproducing both numbers from the component shapes is a strong check that the expert + # layout (non-gated, 2 fused shared experts) and the layer pattern are right. + n_embd, vocab_size = 2688, 131072 + pattern = "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME" + num_mamba = pattern.count("M") + num_moe = pattern.count("E") + num_attn = pattern.count("*") + assert (num_mamba, num_moe, num_attn) == (23, 23, 6) + + d_inner = 64 * 64 + group_state = 8 * 128 + mamba = ( + n_embd * (2 * d_inner + 2 * group_state + 64) + + (d_inner + 2 * group_state) * 4 + + (d_inner + 2 * group_state) + + 64 * 3 + + d_inner + + d_inner * n_embd + ) + routed_experts = 128 * 2 * 1856 * n_embd + shared_experts = 2 * 3712 * n_embd + moe_total = n_embd * 128 + routed_experts + shared_experts + moe_active = n_embd * 128 + 6 * 2 * 1856 * n_embd + shared_experts + attention = (32 * 128 + 2 * 128 + 2 * 128) * n_embd + 32 * 128 * n_embd + embeddings = 2 * vocab_size * n_embd + + total = num_mamba * mamba + num_moe * moe_total + num_attn * attention + embeddings + active = num_mamba * mamba + num_moe * moe_active + num_attn * attention + embeddings + + assert total / 1e9 == pytest.approx(31.6, abs=0.1) + assert active / 1e9 == pytest.approx(3.6, abs=0.1) + # Excluding the input embedding gives the reported 3.2B active parameters. + assert (active - vocab_size * n_embd) / 1e9 == pytest.approx(3.2, abs=0.1) + + +# -------------------------------------------------------------------------------------------- +# Factory and config validation +# -------------------------------------------------------------------------------------------- + + +def test_factory_builds_a_working_model(): + model = NemotronModelFactory.get_nemotron_model(**_model_kwargs()) + out = model({"input_ids": torch.randint(0, VOCAB_SIZE, (1, 4))}) + assert out["logits"].shape == (1, 4, VOCAB_SIZE) + + +def test_factory_supports_meta_device(): + model = NemotronModelFactory.get_nemotron_model(**_model_kwargs(), use_meta_device=True) + assert all(p.is_meta for p in model.parameters()) + # Materializing must yield finite values once reset_parameters has run. + materialized = model.to_empty(device="cpu") + for module in materialized.modules(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + for name, param in materialized.named_parameters(): + if "mixer" in name: + assert torch.isfinite(param).all(), f"{name} is not finite after materialization" + + +def test_factory_rejects_weight_tying_on_meta_device(): + with pytest.raises(ValueError, match="Weight tying is not supported on the meta device"): + NemotronModelFactory.get_nemotron_model(**_model_kwargs(), use_meta_device=True, use_weight_tying=True) + + +def test_config_rejects_layer_count_mismatch(): + with pytest.raises(ValueError, match="does not match the length of layer_pattern"): + NemotronLLMConfig(**_model_kwargs(n_layer=7)) + + +def test_config_rejects_missing_layer_spec(): + specs = _layer_specs() + del specs["*"] + with pytest.raises(ValueError, match=r"layer_pattern uses the layer types \['\*'\]"): + NemotronLLMConfig(**_model_kwargs(layer_specs=specs)) + + +def test_config_rejects_spec_registered_under_the_wrong_symbol(): + specs = _layer_specs() + specs["M"] = specs["-"] + with pytest.raises(ValueError, match="is a spec for layer type"): + NemotronLLMConfig(**_model_kwargs(layer_specs=specs)) + + +def test_config_warns_about_unused_layer_specs(caplog): + with caplog.at_level("WARNING"): + NemotronLLMConfig(**_model_kwargs(layer_pattern="ME", n_layer=2)) + assert "unused layer types" in caplog.text + + +@pytest.mark.parametrize("field,value", [("n_embd", 130), ("vocab_size", 250)]) +def test_config_enforces_tensor_core_alignment(field, value): + with pytest.raises(ValueError, match="should be divisible by 128"): + NemotronLLMConfig(**_model_kwargs(**{field: value})) + + +def test_config_alignment_check_can_be_disabled(): + config = NemotronLLMConfig(**_model_kwargs(n_embd=130, enforce_tensor_core_alignment=False)) + assert config.n_embd == 130 + + +def test_model_rejects_missing_spec_at_construction_time(): + specs = _layer_specs() + del specs["E"] + with pytest.raises(ValueError, match="No layer spec provided"): + _make_model(layer_specs=specs) + + +# -------------------------------------------------------------------------------------------- +# Layer specs +# -------------------------------------------------------------------------------------------- + + +def test_layer_specs_report_their_symbol(): + specs = _layer_specs() + assert specs["M"].symbol == LayerSymbol.MAMBA + assert specs["E"].symbol == LayerSymbol.MOE + assert specs["*"].symbol == LayerSymbol.ATTENTION + assert specs["-"].symbol == LayerSymbol.MLP + + +def test_layer_spec_build_returns_independent_modules(): + spec = _layer_specs()["M"] + first, second = spec.build(layer_idx=0), spec.build(layer_idx=1) + assert first is not second + assert first.mixer.in_proj.weight is not second.mixer.in_proj.weight + + +def test_mamba_spec_rejects_indivisible_group_count(): + with pytest.raises(ValueError, match="must be divisible by"): + Mamba2LayerSpec( + n_embd=N_EMBD, + mamba_n_heads=6, + mamba_head_dim=16, + mamba_state_dim=8, + mamba_n_groups=4, + norm_config=NORM_CONFIG, + ) + + +def test_attention_spec_rejects_indivisible_head_count(): + with pytest.raises(ValueError, match="must be divisible by n_head_kv"): + NemotronAttentionLayerSpec(n_embd=N_EMBD, n_head_q=3, n_head_kv=2, head_dim=16, norm_config=NORM_CONFIG) + + +def test_moe_spec_rejects_top_k_above_num_experts(): + with pytest.raises(ValueError, match="must not exceed num_experts"): + NemotronMoELayerSpec(n_embd=N_EMBD, num_experts=4, moe_ffn_hidden=32, top_k=5, norm_config=NORM_CONFIG) + + +def test_moe_spec_rejects_unknown_router_dtype(): + with pytest.raises(ValueError, match="router_dtype must be"): + NemotronMoELayerSpec( + n_embd=N_EMBD, + num_experts=4, + moe_ffn_hidden=32, + top_k=2, + router_dtype="float64", + norm_config=NORM_CONFIG, + ) + + +def test_moe_spec_fuses_shared_experts_into_one_mlp(): + # The reference implementation represents N shared experts as a single MLP of N times the + # per-expert hidden size (moe_shared_expert_intermediate_size = 2 * 1856 for Nemotron). + spec = NemotronMoELayerSpec( + n_embd=N_EMBD, + num_experts=8, + moe_ffn_hidden=32, + top_k=2, + num_shared_experts=2, + experts_backend=ExpertsBackend.LOOPED, + norm_config=NORM_CONFIG, + ) + assert spec.config.shared_expert_ffn_hidden == 64 + layer = spec.build(layer_idx=0) + assert layer.moe.shared_experts.c_fc.out_features == 64 + + +def test_moe_spec_can_disable_shared_experts(): + spec = NemotronMoELayerSpec( + n_embd=N_EMBD, + num_experts=8, + moe_ffn_hidden=32, + top_k=2, + num_shared_experts=0, + experts_backend=ExpertsBackend.LOOPED, + norm_config=NORM_CONFIG, + ) + assert spec.config.shared_expert_ffn_hidden == 0 + assert spec.build(layer_idx=0).moe.shared_experts is None + + +def test_specs_produce_the_expected_component_types(): + specs = _layer_specs() + assert isinstance(specs["M"].build(0).mixer, Mamba2Mixer) + assert isinstance(specs["E"].build(0).moe, MoE) + assert isinstance(specs["*"].build(0).attn, NemotronSelfAttention) diff --git a/tests/models/nemotron/test_nemotron_stages_and_mfu.py b/tests/models/nemotron/test_nemotron_stages_and_mfu.py new file mode 100644 index 000000000..eed66b2dc --- /dev/null +++ b/tests/models/nemotron/test_nemotron_stages_and_mfu.py @@ -0,0 +1,348 @@ +"""Tests for pipeline stage generation, the MFU calculator and the full 30B-A3B config.""" + +import os +from pathlib import Path + +import pytest +import torch + +from modalities.config.config import load_app_config_dict +from modalities.models.nemotron.layer_pattern import LayerSymbol +from modalities.models.nemotron.nemotron_stages_generator import ( + DEFAULT_LAYER_WEIGHTS, + NemotronStagesGenerator, + NemotronStagesGeneratorConfig, +) +from modalities.utils.nemotron_mfu import NemotronMFUCalculator +from tests.models.nemotron.test_nemotron_model import _make_model + +NEMOTRON_3_NANO_PATTERN = "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME" +FULL_CONFIG_PATH = Path(__file__).parents[3] / "config_files" / "training" / "config_nemotron3_nano_30b_a3b_fsdp2.yaml" + + +def _load_full_config() -> dict: + """Loads the 30B-A3B config, supplying the environment values the resolvers expect. + + The config interpolates the distributed-launch environment (``LOCAL_RANK`` etc.) and the + experiment paths, none of which exist in a plain pytest process. + """ + launch_env = {"LOCAL_RANK": "0", "RANK": "0", "WORLD_SIZE": "1", "LOCAL_WORLD_SIZE": "1"} + previous = {key: os.environ.get(key) for key in launch_env} + os.environ.update(launch_env) + try: + return load_app_config_dict( + config_file_path=FULL_CONFIG_PATH, + experiment_id="test", + experiments_root_path=Path("/tmp/modalities_experiments"), + ) + finally: + for key, value in previous.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +# -------------------------------------------------------------------------------------------- +# Stages generator +# -------------------------------------------------------------------------------------------- + + +def _all_module_names(generator: NemotronStagesGenerator) -> set[str]: + return {fqn for fqns, _ in generator._get_potential_split_points() for fqn in fqns} + + +@pytest.mark.parametrize( + "layer_pattern,num_layers_per_stage,pp_dims", + [ + ("MEM*EM-E", 6, 2), + ("MEM*EM-E", 3, 4), + # Cost-skewed pattern: all expensive layers first. Greedy packing against a fixed cap + # exhausts its budget early here, which is what used to drop the trailing modules. + ("EEEEMMMM", 6, 2), + ("MMMMEEEE", 6, 2), + (NEMOTRON_3_NANO_PATTERN, 14, 4), + (NEMOTRON_3_NANO_PATTERN, 28, 2), + ], +) +def test_stages_cover_every_module_exactly_once(layer_pattern, num_layers_per_stage, pp_dims): + # Regression guard: a partitioner that drops the lm_head produces a model that trains without + # an output layer instead of failing loudly. + generator = NemotronStagesGenerator(layer_pattern=layer_pattern) + stages = generator.get_stages(num_layers_per_stage=num_layers_per_stage, pp_dims=pp_dims) + flat = [fqn for stage in stages for fqn in stage] + + assert set(flat) == _all_module_names(generator), "some module was not assigned to any stage" + assert len(flat) == len(set(flat)), "a module was assigned to more than one stage" + assert all(stage for stage in stages), "an empty pipeline stage was produced" + assert len(stages) % pp_dims == 0 + + +def test_stages_include_the_language_model_head(): + generator = NemotronStagesGenerator(layer_pattern="EEEEMMMM") + stages = generator.get_stages(num_layers_per_stage=6, pp_dims=2) + assert "transformer.lm_head" in stages[-1] + assert "transformer.lm_head_norm" in stages[-1] + assert "transformer.wte" in stages[0] + + +def test_stages_preserve_model_order(): + generator = NemotronStagesGenerator(layer_pattern="MEMEM*E-") + stages = generator.get_stages(num_layers_per_stage=6, pp_dims=2) + flat = [fqn for stage in stages for fqn in stage] + layer_indices = [int(fqn.split(".")[-1]) for fqn in flat if fqn.startswith("transformer.h.")] + assert layer_indices == sorted(layer_indices) + + +def test_stage_weights_are_near_balanced_for_the_full_model(): + generator = NemotronStagesGenerator(layer_pattern=NEMOTRON_3_NANO_PATTERN) + split_points = generator._get_potential_split_points() + weight_by_fqn = {fqns[0]: weight for fqns, weight in split_points} + stages = generator.get_stages(num_layers_per_stage=14, pp_dims=4) + + stage_weights = [sum(weight_by_fqn[fqn] for fqn in stage if fqn in weight_by_fqn) for stage in stages] + ideal = sum(weight_by_fqn.values()) / len(stages) + # The slowest stage sets pipeline throughput, so the imbalance must stay small. + assert max(stage_weights) <= ideal * 1.1 + + +def test_stages_generator_rejects_more_stages_than_modules(): + generator = NemotronStagesGenerator(layer_pattern="ME") + with pytest.raises(ValueError, match="Cannot build"): + generator.get_stages(num_layers_per_stage=1, pp_dims=6) + + +def test_stages_generator_rejects_non_divisible_stage_counts(): + generator = NemotronStagesGenerator(layer_pattern="MEM*EM-E") + with pytest.raises(ValueError, match="not divisible by parallel dimensions"): + generator.get_stages(num_layers_per_stage=4, pp_dims=2) + + +def test_stages_generator_rejects_invalid_layers_per_stage(): + generator = NemotronStagesGenerator(layer_pattern="ME") + with pytest.raises(ValueError, match="num_layers_per_stage must be at least 1"): + generator.get_stages(num_layers_per_stage=0, pp_dims=1) + + +def test_moe_layers_are_weighted_more_heavily_than_mamba_layers(): + # The point of this generator: equal *counts* of layers would give unbalanced stages, because + # an MoE layer performs far more work per token than a Mamba layer. + assert DEFAULT_LAYER_WEIGHTS[LayerSymbol.MOE] > DEFAULT_LAYER_WEIGHTS[LayerSymbol.MAMBA] + assert DEFAULT_LAYER_WEIGHTS[LayerSymbol.MAMBA] > DEFAULT_LAYER_WEIGHTS[LayerSymbol.ATTENTION] + + +def test_weighted_split_balances_cost_rather_than_layer_count(): + # An "all MoE first, all Mamba second" pattern: a count-based split would put 4 layers in each + # stage, but the MoE half costs 50% more. The weighted split must give the MoE half fewer + # layers so that the two stages take a comparable amount of time. + generator = NemotronStagesGenerator(layer_pattern="EEEEMMMM") + stages = generator.get_stages(num_layers_per_stage=6, pp_dims=2) + first_stage_layers = [f for f in stages[0] if f.startswith("transformer.h.")] + last_stage_layers = [f for f in stages[-1] if f.startswith("transformer.h.")] + assert len(first_stage_layers) < len(last_stage_layers) + + +def test_custom_layer_weights_are_honoured(): + generator = NemotronStagesGenerator( + layer_pattern="ME", + layer_weights={"M": 1, "E": 1, "*": 1, "-": 1}, + ) + split_points = generator._get_potential_split_points() + layer_weights = [weight for fqns, weight in split_points if fqns[0].startswith("transformer.h.")] + assert layer_weights == [1, 1] + + +def test_default_layer_weights_are_applied_per_layer_type(): + generator = NemotronStagesGenerator(layer_pattern="ME*-") + split_points = generator._get_potential_split_points() + layer_weights = [weight for fqns, weight in split_points if fqns[0].startswith("transformer.h.")] + assert layer_weights == [ + DEFAULT_LAYER_WEIGHTS[LayerSymbol.MAMBA], + DEFAULT_LAYER_WEIGHTS[LayerSymbol.MOE], + DEFAULT_LAYER_WEIGHTS[LayerSymbol.ATTENTION], + DEFAULT_LAYER_WEIGHTS[LayerSymbol.MLP], + ] + + +def test_stage_fqns_resolve_against_a_real_model(): + model = _make_model(layer_pattern="ME*-", n_layer=4) + generator = NemotronStagesGenerator(layer_pattern="ME*-") + stages = generator.get_stages(num_layers_per_stage=4, pp_dims=2) + for stage in stages: + for fqn in stage: + assert model.get_submodule(fqn) is not None, fqn + + +def test_stages_generator_config_validates_the_pattern(): + with pytest.raises(ValueError, match="Invalid layer symbol"): + NemotronStagesGeneratorConfig(layer_pattern="MEX") + + +def test_stages_generator_config_requires_weights_for_every_used_layer_type(): + with pytest.raises(ValueError, match=r"missing an entry for the layer types \['E'\]"): + NemotronStagesGeneratorConfig(layer_pattern="ME", layer_weights={"M": 1}) + + +def test_nemotron_3_nano_splits_into_four_pipeline_stages(): + generator = NemotronStagesGenerator(layer_pattern=NEMOTRON_3_NANO_PATTERN) + total_weight = sum(weight for _, weight in generator._get_potential_split_points()) + stages = generator.get_stages(num_layers_per_stage=14, pp_dims=4) + assert len(stages) % 4 == 0 + flat = [fqn for stage in stages for fqn in stage] + assert len(flat) == 52 + 3 # 52 layers + wte + lm_head_norm + lm_head + assert total_weight == 23 * 2 + 23 * 3 + 6 * 1 + 2 + 2 # mamba + moe + attention + in + out + + +# -------------------------------------------------------------------------------------------- +# MFU calculator +# -------------------------------------------------------------------------------------------- + + +def test_active_parameter_count_excludes_unrouted_experts(): + model = _make_model(layer_pattern="E", n_layer=1) + total = sum(p.numel() for p in model.parameters()) + active = NemotronMFUCalculator.count_active_parameters(model) + + moe = model.transformer.h["0"].moe + routed = sum(p.numel() for p in moe.experts.parameters()) + expected = total - int(routed * (1 - moe.router.top_k / moe.router.num_experts)) + assert active == expected + # top-2 of 8 experts: the routed experts are the dominant parameter block, so the active + # count must be substantially smaller than the total. + assert active < total + + +def test_active_parameter_count_equals_total_for_a_dense_model(): + model = _make_model(layer_pattern="M-", n_layer=2) + total = sum(p.numel() for p in model.parameters()) + assert NemotronMFUCalculator.count_active_parameters(model) == total + + +def test_nemotron_3_nano_active_parameter_ratio(): + # Model report: 31.6B total, 3.6B active including embeddings -> roughly a 9x sparsity factor. + total, active = 31.6e9, 3.6e9 + assert total / active == pytest.approx(8.8, abs=0.3) + + +def test_flops_per_token_counts_only_attention_layers_quadratically(): + # 6 of 52 layers attend; a formula that charged the quadratic term to all 52 would inflate the + # theoretical FLOPs by roughly 8x and report a correspondingly deflated MFU. + common = dict(num_active_params=1_000_000, sequence_length=8192, n_head_q=32, head_dim=128) + six_layers = NemotronMFUCalculator._get_theoretical_flops_per_token(num_attention_layers=6, **common) + all_layers = NemotronMFUCalculator._get_theoretical_flops_per_token(num_attention_layers=52, **common) + assert all_layers > six_layers + + dense_only = NemotronMFUCalculator._get_theoretical_flops_per_token(num_attention_layers=0, **common) + assert dense_only == 6 * 1_000_000 + assert six_layers - dense_only == 12 * 6 * 8192 * 32 * 128 + + +def test_flops_per_token_scales_linearly_in_active_parameters(): + common = dict(num_attention_layers=6, sequence_length=1024, n_head_q=8, head_dim=64) + small = NemotronMFUCalculator._get_theoretical_flops_per_token(num_active_params=1_000, **common) + large = NemotronMFUCalculator._get_theoretical_flops_per_token(num_active_params=2_000, **common) + assert large - small == 6 * 1_000 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="peak performance lookup requires a GPU") +def test_mfu_calculator_produces_a_finite_utilization(): + model = _make_model(layer_pattern="ME*-", n_layer=4).cuda() + calculator = NemotronMFUCalculator( + layer_pattern="ME*-", + sequence_length=32, + n_embd=model.n_embd, + n_head_q=4, + head_dim=16, + num_active_params=NemotronMFUCalculator.count_active_parameters(model), + world_size=1, + # The shared peak-performance lookup expects an FSDP-wrapped model or a list of pipeline + # stages; a single-element list is the unwrapped equivalent. + model_parts=[model], + ) + mfu = calculator.compute(num_samples_per_second=torch.tensor(10.0)) + assert torch.isfinite(mfu) + assert mfu > 0 + + +# -------------------------------------------------------------------------------------------- +# The full 30B-A3B configuration +# -------------------------------------------------------------------------------------------- + + +def test_full_config_file_is_parseable(): + assert FULL_CONFIG_PATH.exists(), FULL_CONFIG_PATH + config_dict = _load_full_config() + assert config_dict["model_raw"]["variant_key"] == "nemotron" + + +def test_full_config_matches_the_published_architecture(): + config = _load_full_config()["model_raw"]["config"] + assert config["n_layer"] == 52 + assert config["n_embd"] == 2688 + assert config["layer_pattern"] == NEMOTRON_3_NANO_PATTERN + assert config["vocab_size"] == 131072 + assert config["use_weight_tying"] is False + + mamba = config["layer_specs"]["M"]["config"] + assert (mamba["mamba_n_heads"], mamba["mamba_head_dim"]) == (64, 64) + assert (mamba["mamba_state_dim"], mamba["mamba_n_groups"]) == (128, 8) + assert mamba["d_conv"] == 4 and mamba["chunk_size"] == 128 + + moe = config["layer_specs"]["E"]["config"] + assert moe["num_experts"] == 128 + assert moe["moe_ffn_hidden"] == 1856 + assert moe["top_k"] == 6 + assert moe["num_shared_experts"] == 2 + assert moe["shared_expert_ffn_hidden_per_expert"] == 1856 + assert moe["score_function"] == "sigmoid" + assert moe["use_expert_bias"] is True + assert moe["route_scale"] == 2.5 + assert moe["aux_loss_coeff"] == pytest.approx(1e-4) + + attention = config["layer_specs"]["*"]["config"] + assert (attention["n_head_q"], attention["n_head_kv"], attention["head_dim"]) == (32, 2, 128) + + +def test_full_config_wires_load_balancing_and_the_auxiliary_loss(): + config_dict = _load_full_config() + + optimizer = config_dict["optimizer"] + assert optimizer["variant_key"] == "moe_load_balanced" + assert optimizer["config"]["expert_bias_update_rate"] == pytest.approx(1e-3) + excluded = optimizer["config"]["optimizer"]["config"]["weight_decay_groups_excluded"] + assert set(excluded) == {"embedding", "layernorm", "ssm", "router"} + + loss = config_dict["loss_fn"] + assert loss["variant_key"] == "weighted_sum" + variants = [term["variant_key"] for term in loss["config"]["losses"]] + assert variants == ["clm_cross_entropy_loss", "moe_aux_loss"] + assert config_dict["model_raw"]["config"]["aux_loss_key"] == "moe_aux_loss" + + +def test_full_config_declares_every_layer_class_as_an_fsdp_block(): + config_dict = _load_full_config() + block_names = set(config_dict["fsdp_model"]["config"]["block_names"]) + assert block_names == {"Mamba2Layer", "NemotronMoELayer", "NemotronAttentionLayer", "NemotronMLPLayer"} + assert config_dict["activation_checkpointed_model"]["config"]["layers_fqn"] == "transformer.h" + + +def test_full_config_layer_specs_validate(): + # Build the layer specs from the real config to confirm the published hyperparameters satisfy + # every constraint (head/group divisibility, grouped_mm alignment, top_k <= num_experts). + from modalities.models.nemotron.nemotron_layer_specs import ( + Mamba2LayerSpec, + NemotronAttentionLayerSpec, + NemotronMoELayerSpec, + ) + + config = _load_full_config()["model_raw"]["config"] + norm_config = {"norm_type": "pytorch_rms_norm", "config": {"normalized_shape": 2688, "eps": 1e-5}} + + mamba_spec = Mamba2LayerSpec(**{**config["layer_specs"]["M"]["config"], "norm_config": norm_config}) + assert mamba_spec.config.mamba_n_heads * mamba_spec.config.mamba_head_dim == 4096 + + moe_spec = NemotronMoELayerSpec(**{**config["layer_specs"]["E"]["config"], "norm_config": norm_config}) + assert moe_spec.config.shared_expert_ffn_hidden == 3712 + + attention_spec = NemotronAttentionLayerSpec(**{**config["layer_specs"]["*"]["config"], "norm_config": norm_config}) + assert attention_spec.config.n_head_q // attention_spec.config.n_head_kv == 16 diff --git a/tests/models/nemotron/test_ssd.py b/tests/models/nemotron/test_ssd.py new file mode 100644 index 000000000..f2c2aa56b --- /dev/null +++ b/tests/models/nemotron/test_ssd.py @@ -0,0 +1,267 @@ +"""Correctness tests for the pure-PyTorch Mamba-2 SSD primitives. + +The chunked scan is the numerical foundation of every Mamba layer, so it is validated against a +literal step-by-step transcription of the recurrence (:func:`ssd_recurrent_reference`) rather than +against itself. +""" + +import pytest +import torch + +from modalities.models.components.mamba2.ssd import ( + GatedRMSNorm, + causal_depthwise_conv1d, + ssd_chunked_scan, + ssd_recurrent_reference, +) + + +def _make_ssd_inputs( + batch_size: int = 2, + seq_len: int = 24, + num_heads: int = 4, + head_dim: int = 8, + num_groups: int = 2, + state_dim: int = 6, + seed: int = 0, +) -> dict[str, torch.Tensor]: + """Builds a small, deterministic set of SSD inputs in float32.""" + generator = torch.Generator().manual_seed(seed) + + def randn(*shape): + return torch.randn(*shape, generator=generator, dtype=torch.float32) + + return { + "x": randn(batch_size, seq_len, num_heads, head_dim), + # dt must be strictly positive; softplus of a standard normal is a realistic range. + "dt": torch.nn.functional.softplus(randn(batch_size, seq_len, num_heads)), + # A is strictly negative, i.e. -exp(A_log) with A_log = log(U(1, 16)). + "A": -torch.empty(num_heads).uniform_(1.0, 16.0, generator=generator), + "B": randn(batch_size, seq_len, num_groups, state_dim), + "C": randn(batch_size, seq_len, num_groups, state_dim), + "D": torch.ones(num_heads), + } + + +@pytest.mark.parametrize("chunk_size", [1, 2, 4, 8, 24]) +def test_chunked_scan_matches_recurrent_reference(chunk_size): + inputs = _make_ssd_inputs() + expected = ssd_recurrent_reference(**inputs) + actual = ssd_chunked_scan(**inputs, chunk_size=chunk_size) + torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-4) + + +@pytest.mark.parametrize("chunk_size", [4, 8, 16]) +def test_chunked_scan_is_invariant_to_chunk_size(chunk_size): + # The block decomposition is an implementation detail: the mathematical result must not + # depend on how the sequence is split. + inputs = _make_ssd_inputs(seq_len=32) + reference = ssd_chunked_scan(**inputs, chunk_size=32) + actual = ssd_chunked_scan(**inputs, chunk_size=chunk_size) + torch.testing.assert_close(actual, reference, rtol=1e-4, atol=1e-4) + + +@pytest.mark.parametrize("seq_len", [1, 5, 7, 13]) +def test_chunked_scan_handles_sequences_not_divisible_by_chunk_size(seq_len): + inputs = _make_ssd_inputs(seq_len=seq_len) + expected = ssd_recurrent_reference(**inputs) + actual = ssd_chunked_scan(**inputs, chunk_size=4) + assert actual.shape == expected.shape + torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-4) + + +def test_chunked_scan_is_causal(): + # Perturbing token t must leave every output at position < t untouched. + inputs = _make_ssd_inputs(seq_len=16) + baseline = ssd_chunked_scan(**inputs, chunk_size=4) + + perturbed = {key: value.clone() for key, value in inputs.items()} + perturb_at = 9 + perturbed["x"][:, perturb_at] += 5.0 + outputs = ssd_chunked_scan(**perturbed, chunk_size=4) + + torch.testing.assert_close(outputs[:, :perturb_at], baseline[:, :perturb_at], rtol=1e-5, atol=1e-6) + assert not torch.allclose(outputs[:, perturb_at], baseline[:, perturb_at]) + + +def test_chunked_scan_zero_input_gives_zero_output_without_skip(): + inputs = _make_ssd_inputs() + inputs["x"] = torch.zeros_like(inputs["x"]) + inputs["D"] = None + outputs = ssd_chunked_scan(**inputs, chunk_size=4) + torch.testing.assert_close(outputs, torch.zeros_like(outputs)) + + +def test_chunked_scan_skip_connection_is_additive(): + inputs = _make_ssd_inputs() + without_skip = ssd_chunked_scan(**{**inputs, "D": None}, chunk_size=4) + scale = 0.5 + with_skip = ssd_chunked_scan(**{**inputs, "D": torch.full_like(inputs["D"], scale)}, chunk_size=4) + expected = without_skip + scale * inputs["x"] + torch.testing.assert_close(with_skip, expected, rtol=1e-4, atol=1e-5) + + +def test_chunked_scan_respects_initial_state(): + inputs = _make_ssd_inputs(seq_len=8) + batch_size, _, num_heads, head_dim = inputs["x"].shape + state_dim = inputs["B"].shape[-1] + zero_state = torch.zeros(batch_size, num_heads, head_dim, state_dim) + + with_zero_state = ssd_chunked_scan(**inputs, chunk_size=4, initial_states=zero_state) + without_state = ssd_chunked_scan(**inputs, chunk_size=4) + torch.testing.assert_close(with_zero_state, without_state, rtol=1e-5, atol=1e-6) + + nonzero_state = torch.ones_like(zero_state) + with_nonzero_state = ssd_chunked_scan(**inputs, chunk_size=4, initial_states=nonzero_state) + assert not torch.allclose(with_nonzero_state, without_state) + + +def test_chunked_scan_final_state_matches_split_sequence_run(): + # Running the full sequence must equal running the first half, then the second half seeded + # with the returned state. This pins down the inter-chunk state passing. + inputs = _make_ssd_inputs(seq_len=16) + full = ssd_chunked_scan(**inputs, chunk_size=4) + + def slice_inputs(start: int, stop: int) -> dict[str, torch.Tensor]: + sliced = {key: value[:, start:stop] for key, value in inputs.items() if key in ("x", "dt", "B", "C")} + sliced["A"] = inputs["A"] + sliced["D"] = inputs["D"] + return sliced + + first_half, state = ssd_chunked_scan(**slice_inputs(0, 8), chunk_size=4, return_final_state=True) + second_half = ssd_chunked_scan(**slice_inputs(8, 16), chunk_size=4, initial_states=state) + + torch.testing.assert_close(torch.cat([first_half, second_half], dim=1), full, rtol=1e-4, atol=1e-4) + + +def test_chunked_scan_shares_bc_across_heads_of_a_group(): + # With one group, all heads must see the same B/C, which is what makes the group expansion + # equivalent to explicitly broadcasting B/C. + inputs = _make_ssd_inputs(num_heads=4, num_groups=1) + grouped = ssd_chunked_scan(**inputs, chunk_size=4) + + expanded = dict(inputs) + expanded["B"] = inputs["B"].expand(-1, -1, 4, -1).contiguous() + expanded["C"] = inputs["C"].expand(-1, -1, 4, -1).contiguous() + # With B/C already per-head, the expansion inside the scan is a no-op. + per_head = ssd_chunked_scan(**expanded, chunk_size=4) + torch.testing.assert_close(grouped, per_head, rtol=1e-5, atol=1e-6) + + +def test_chunked_scan_rejects_invalid_arguments(): + inputs = _make_ssd_inputs() + with pytest.raises(ValueError, match="chunk_size must be positive"): + ssd_chunked_scan(**inputs, chunk_size=0) + with pytest.raises(ValueError, match=r"x must have shape"): + ssd_chunked_scan(**{**inputs, "x": inputs["x"][:, :, 0]}, chunk_size=4) + with pytest.raises(ValueError, match="must be divisible by num_groups"): + ssd_chunked_scan(**{**inputs, "x": inputs["x"][:, :, :3]}, chunk_size=4) + + +def test_chunked_scan_preserves_input_dtype(): + inputs = _make_ssd_inputs() + bf16_inputs = {key: (value.bfloat16() if key != "A" else value) for key, value in inputs.items()} + outputs = ssd_chunked_scan(**bf16_inputs, chunk_size=4) + assert outputs.dtype == torch.bfloat16 + + +def test_chunked_scan_is_differentiable(): + inputs = _make_ssd_inputs(seq_len=8) + inputs["x"].requires_grad_(True) + outputs = ssd_chunked_scan(**inputs, chunk_size=4) + outputs.sum().backward() + assert inputs["x"].grad is not None + assert torch.isfinite(inputs["x"].grad).all() + + +def test_causal_depthwise_conv1d_is_causal_and_matches_manual_convolution(): + torch.manual_seed(0) + batch_size, seq_len, channels, kernel_size = 2, 7, 5, 4 + x = torch.randn(batch_size, seq_len, channels) + weight = torch.randn(channels, 1, kernel_size) + bias = torch.randn(channels) + + actual = causal_depthwise_conv1d(x, weight=weight, bias=bias) + assert actual.shape == x.shape + + # Manual reference: y[t, c] = bias[c] + sum_k w[c, 0, k] * x[t - (K - 1) + k, c] + expected = torch.zeros_like(actual) + for t in range(seq_len): + for k in range(kernel_size): + source = t - (kernel_size - 1) + k + if source >= 0: + expected[:, t] += weight[:, 0, k] * x[:, source] + expected += bias + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6) + + +def test_causal_depthwise_conv1d_does_not_leak_future_tokens(): + torch.manual_seed(0) + x = torch.randn(1, 10, 3) + weight = torch.randn(3, 1, 4) + baseline = causal_depthwise_conv1d(x, weight=weight, bias=None) + + perturbed = x.clone() + perturbed[:, 6] += 10.0 + outputs = causal_depthwise_conv1d(perturbed, weight=weight, bias=None) + torch.testing.assert_close(outputs[:, :6], baseline[:, :6], rtol=1e-5, atol=1e-6) + + +def test_causal_depthwise_conv1d_is_channelwise_independent(): + torch.manual_seed(0) + x = torch.randn(1, 6, 4) + weight = torch.randn(4, 1, 3) + outputs = causal_depthwise_conv1d(x, weight=weight, bias=None) + + perturbed = x.clone() + perturbed[:, :, 2] += 3.0 + perturbed_outputs = causal_depthwise_conv1d(perturbed, weight=weight, bias=None) + unchanged = [0, 1, 3] + torch.testing.assert_close(perturbed_outputs[:, :, unchanged], outputs[:, :, unchanged], rtol=1e-5, atol=1e-6) + + +def test_gated_rms_norm_normalizes_per_group(): + torch.manual_seed(0) + hidden_size, num_groups = 8, 2 + norm = GatedRMSNorm(hidden_size=hidden_size, num_groups=num_groups) + x = torch.randn(3, 5, hidden_size) + # A constant gate contributes a constant positive factor, so the per-group root mean square + # after normalization must be exactly 1 (the learnable weight is initialized to 1). + gate = torch.full_like(x, 10.0) + out = norm(x, gate=gate) + + grouped = out.reshape(3, 5, num_groups, hidden_size // num_groups) + rms = grouped.pow(2).mean(dim=-1).sqrt() + torch.testing.assert_close(rms, torch.ones_like(rms), rtol=2e-3, atol=2e-3) + + +def test_gated_rms_norm_groups_are_normalized_independently(): + # Scaling only the second group must leave the first group's output untouched: that is the + # difference between grouped and global RMS normalization. + torch.manual_seed(0) + norm = GatedRMSNorm(hidden_size=8, num_groups=2) + x = torch.randn(2, 3, 8) + gate = torch.randn(2, 3, 8) + baseline = norm(x, gate=gate) + + scaled = x.clone() + scaled[..., 4:] *= 100.0 + outputs = norm(scaled, gate=gate) + torch.testing.assert_close(outputs[..., :4], baseline[..., :4], rtol=1e-5, atol=1e-6) + + +def test_gated_rms_norm_applies_gate_before_normalization(): + torch.manual_seed(0) + norm = GatedRMSNorm(hidden_size=4, num_groups=1) + x = torch.randn(2, 3, 4) + gate = torch.randn(2, 3, 4) + + expected_pre_norm = x.float() * torch.nn.functional.silu(gate.float()) + inv_rms = torch.rsqrt(expected_pre_norm.pow(2).mean(dim=-1, keepdim=True) + norm.eps) + expected = expected_pre_norm * inv_rms + torch.testing.assert_close(norm(x, gate=gate), expected, rtol=1e-5, atol=1e-6) + + +def test_gated_rms_norm_rejects_indivisible_group_count(): + with pytest.raises(ValueError, match="must be divisible by num_groups"): + GatedRMSNorm(hidden_size=7, num_groups=2) diff --git a/tests/test_yaml_configs/nemotron_config_initialization.yaml b/tests/test_yaml_configs/nemotron_config_initialization.yaml new file mode 100644 index 000000000..4ff0567e6 --- /dev/null +++ b/tests/test_yaml_configs/nemotron_config_initialization.yaml @@ -0,0 +1,84 @@ +model: + component_key: model + variant_key: model_initialized + config: + model: + instance_key: model_raw + pass_type: BY_REFERENCE + model_initializer: + component_key: model_initialization + variant_key: composed + config: + model_type: nemotron + weight_init_type: WILL_BE_REPLACED + mean: 0.0 + std: WILL_BE_REPLACED + num_layers: ${model_raw.config.n_layer} + +model_raw: + component_key: model + variant_key: nemotron + config: + sample_key: "input_ids" + prediction_key: "logits" + aux_loss_key: "moe_aux_loss" + sequence_length: 256 + vocab_size: 512 + n_embd: 256 + n_layer: 8 + # A miniature of the Nemotron-3 Nano pattern: mostly Mamba/MoE with one attention layer and + # one dense MLP layer, so that all four layer types are exercised. + layer_pattern: "MEMEM*E-" + use_weight_tying: false + lm_head_norm_config: &nemotron_norm_config + norm_type: pytorch_rms_norm + config: + normalized_shape: ${model_raw.config.n_embd} + eps: 1e-5 + layer_specs: + "M": + component_key: nemotron_layer_spec + variant_key: mamba2 + config: + n_embd: ${model_raw.config.n_embd} + mamba_n_heads: 8 + mamba_head_dim: 32 + mamba_state_dim: 16 + mamba_n_groups: 2 + d_conv: 4 + chunk_size: 32 + ssd_backend: native + norm_config: *nemotron_norm_config + "E": + component_key: nemotron_layer_spec + variant_key: moe + config: + n_embd: ${model_raw.config.n_embd} + num_experts: 8 + moe_ffn_hidden: 64 + top_k: 2 + route_scale: 2.5 + score_function: sigmoid + use_expert_bias: true + router_dtype: float32 + num_shared_experts: 2 + aux_loss_coeff: 1.0e-4 + experts_backend: looped + norm_config: *nemotron_norm_config + "*": + component_key: nemotron_layer_spec + variant_key: attention + config: + n_embd: ${model_raw.config.n_embd} + n_head_q: 8 + n_head_kv: 2 + head_dim: 32 + attention_implementation: pytorch_flash + norm_config: *nemotron_norm_config + "-": + component_key: nemotron_layer_spec + variant_key: mlp + config: + n_embd: ${model_raw.config.n_embd} + ffn_hidden: 128 + norm_config: *nemotron_norm_config From b646a67a4c5fa65092975d7bfa6bc6db44320e6e Mon Sep 17 00:00:00 2001 From: Max Luebbering <2804731+le1nux@users.noreply.github.com> Date: Tue, 28 Jul 2026 18:51:06 +0200 Subject: [PATCH 2/4] docs(model): attribute adopted code in the Nemotron implementation Modalities is MIT licensed; Megatron-LM, Megatron-Bridge and state-spaces/mamba are Apache 2.0 (NVIDIA / Tri Dao, Albert Gu) and TorchTitan is BSD 3-Clause. Preserving those notices in the derived files is a license obligation, so each source file that adapts upstream code now carries a header naming the upstream file, what was taken from it, the copyright holder and the license. Marked as adapted from Megatron-LM: * mamba2_mixer.py - packed [z, x, B, C, dt] projection layout, conv1d/A_log/D/ dt_bias shapes and initialization distributions * layer_pattern.py - the layer pattern symbols * nemotron_layers.py - single-operator pre-norm residual layer structure * nemotron_layer_specs.py - the declarative ModuleSpec/builder pattern * nemotron_model.py - hybrid model structure * router.py - sigmoid scoring, selection-only expert bias, renormalization * moe.py - the load-balancing loss formula and its sequence-level variant * load_balancing.py - the sign-based expert bias update rule * experts.py, nemotron_mlp.py - the squared ReLU activation Marked as adapted from state-spaces/mamba: * ssd.py - the chunk-parallel SSD block decomposition (ssd_minimal_discrete / segsum) and the GatedRMSNorm semantics Marked as adapted from TorchTitan (following the convention already used in model_factory.py and trainer.py): * experts.py - stacked per-expert weights and the torch._grouped_mm call * moe.py - dispatch/combine structure, expert-bias and token-count buffers * router.py - router interface shape * load_balancing.py - the optimizer step pre-hook approach * nemotron_stages_generator.py - pipeline split-point structure Marked as adapted from Meta's Llama: * nemotron_attention.py - grouped-query key/value head repetition The YAML configs note that their hyperparameter values are adopted from the Megatron-Bridge recipe; no code was taken from Megatron-Bridge. docs/components/nemotron.md gains an Attribution section tabulating all of the above in one place, and explicitly lists the files with no third-party derivation (norms.py, nemotron_model_factory.py, nemotron_mfu.py, and the partitioning algorithm in nemotron_stages_generator.py). Co-Authored-By: Claude Opus 5 (1M context) --- .../config_nemotron3_nano_30b_a3b_fsdp2.yaml | 4 +- docs/components/nemotron.md | 55 +++++++++++++++++++ .../models/components/mamba2/mamba2_mixer.py | 6 ++ .../models/components/mamba2/ssd.py | 6 ++ .../models/components/moe/experts.py | 8 +++ .../models/components/moe/load_balancing.py | 8 +++ src/modalities/models/components/moe/moe.py | 9 +++ .../models/components/moe/router.py | 8 +++ .../models/nemotron/layer_pattern.py | 5 ++ .../models/nemotron/nemotron_attention.py | 4 ++ .../models/nemotron/nemotron_layer_specs.py | 5 ++ .../models/nemotron/nemotron_layers.py | 6 ++ .../models/nemotron/nemotron_mlp.py | 3 + .../models/nemotron/nemotron_model.py | 8 +++ .../nemotron/nemotron_stages_generator.py | 4 ++ .../nemotron_fsdp2_config.yaml | 4 ++ tests/models/nemotron/test_ssd.py | 4 ++ .../nemotron_config_initialization.yaml | 4 ++ 18 files changed, 150 insertions(+), 1 deletion(-) diff --git a/config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml b/config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml index 8a9186d7c..80614905a 100644 --- a/config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml +++ b/config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml @@ -1,7 +1,9 @@ # Nemotron-3 Nano 30B-A3B pretraining configuration. # # Architecture from the model report (arXiv:2512.20848, Table 1 / Figure 2) and cross-checked -# against the Megatron-Bridge recipe `recipes/nemotronh/h100/nemotron_3_nano.py`: +# against the Megatron-Bridge recipe `src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py` +# (Copyright (c) 2026, NVIDIA CORPORATION, licensed under the Apache License, Version 2.0), from +# which the hyperparameter values below are adopted: # # 52 layers pattern MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME # -> 23 Mamba-2, 23 MoE, 6 attention diff --git a/docs/components/nemotron.md b/docs/components/nemotron.md index b63e2cffe..d114bb13b 100644 --- a/docs/components/nemotron.md +++ b/docs/components/nemotron.md @@ -177,3 +177,58 @@ Decaying the SSM dynamics parameters or the router gate destabilizes training. | `stages_generator` | `nemotron_stages_generator` | [NemotronStagesGenerator](../../src/modalities/models/nemotron/nemotron_stages_generator.py) | | `mfu_calculator` | `nemotron` | [NemotronMFUCalculator](../../src/modalities/utils/nemotron_mfu.py) | | `model_initialization` | `composed` (`model_type: nemotron`) | [ComposedInitializationRoutines](../../src/modalities/nn/model_initialization/composed_initialization.py) | + +## Attribution + +Modalities is MIT licensed. The files below adapt code from third-party projects and carry the +corresponding notices in their file headers, as required by those licenses. + +**NVIDIA Megatron-LM** — Copyright (c) NVIDIA CORPORATION, Apache License 2.0. +Portions additionally Copyright (c) 2024, Tri Dao, Albert Gu. + +| File | What was adapted | Upstream | +|------|------------------|----------| +| [mamba2_mixer.py](../../src/modalities/models/components/mamba2/mamba2_mixer.py) | Packed `[z, x, B, C, dt]` projection layout; `conv1d`/`A_log`/`D`/`dt_bias` shapes and init distributions | `megatron/core/ssm/mamba_mixer.py` | +| [layer_pattern.py](../../src/modalities/models/nemotron/layer_pattern.py) | Layer pattern symbols (`M`, `*`, `E`, `-`) | `megatron/core/models/hybrid/hybrid_layer_allocation.py::Symbols` | +| [nemotron_layers.py](../../src/modalities/models/nemotron/nemotron_layers.py) | Single-operator pre-norm residual layer structure | `megatron/core/ssm/mamba_layer.py`, `.../hybrid/hybrid_block.py` | +| [nemotron_layer_specs.py](../../src/modalities/models/nemotron/nemotron_layer_specs.py) | Declarative layer-spec/builder pattern | `megatron/core/transformer/spec_utils.py::ModuleSpec`, `.../hybrid/hybrid_layer_specs.py` | +| [nemotron_model.py](../../src/modalities/models/nemotron/nemotron_model.py) | Hybrid model structure (pattern-driven stack) | `megatron/core/models/hybrid/hybrid_model.py` | +| [router.py](../../src/modalities/models/components/moe/router.py) | Sigmoid scoring, selection-only expert bias, top-k renormalization | `megatron/core/transformer/moe/moe_utils.py::topk_routing_with_score_function` | +| [moe.py](../../src/modalities/models/components/moe/moe.py) | Load-balancing loss formula and its sequence-level variant | `.../moe_utils.py::switch_load_balancing_loss_func` | +| [load_balancing.py](../../src/modalities/models/components/moe/load_balancing.py) | Sign-based expert bias update rule | `.../moe_utils.py::get_updated_expert_bias` | +| [experts.py](../../src/modalities/models/components/moe/experts.py), [nemotron_mlp.py](../../src/modalities/models/nemotron/nemotron_mlp.py) | Squared ReLU activation | `megatron/core/activations.py` | + +**state-spaces/mamba** — Copyright (c) 2024, Tri Dao, Albert Gu, Apache License 2.0. + +| File | What was adapted | Upstream | +|------|------------------|----------| +| [ssd.py](../../src/modalities/models/components/mamba2/ssd.py) | Chunk-parallel SSD block decomposition (`ssd_minimal_discrete` / `segsum`); `GatedRMSNorm` semantics | `ssd_minimal_discrete`, `mamba_ssm.ops.triton.layernorm_gated.rmsnorm_fn` | + +**Meta TorchTitan** — BSD 3-Clause License. (Modalities already adapts from TorchTitan elsewhere.) + +| File | What was adapted | Upstream | +|------|------------------|----------| +| [experts.py](../../src/modalities/models/components/moe/experts.py) | Stacked per-expert weight layout; `torch._grouped_mm` over expert-sorted tokens | `torchtitan/models/common/moe.py::GroupedExperts` | +| [moe.py](../../src/modalities/models/components/moe/moe.py) | Dispatch/combine structure; expert-bias and token-count buffers | `.../moe.py::MoE` | +| [router.py](../../src/modalities/models/components/moe/router.py) | Router interface shape | `.../moe.py::TokenChoiceTopKRouter` | +| [load_balancing.py](../../src/modalities/models/components/moe/load_balancing.py) | Applying the bias update as an optimizer step pre-hook | `.../moe.py` | +| [nemotron_stages_generator.py](../../src/modalities/models/nemotron/nemotron_stages_generator.py) | Pipeline split-point structure (via Modalities' own `StagesGenerator`) | TorchTitan pipeline utilities | + +**Meta Llama** — [facebookresearch/llama](https://github.com/facebookresearch/llama). + +| File | What was adapted | Upstream | +|------|------------------|----------| +| [nemotron_attention.py](../../src/modalities/models/nemotron/nemotron_attention.py) | Grouped-query key/value head repetition (`_repeat_kv`), via Modalities' own GPT-2 | `llama/model.py` | + +**NVIDIA Megatron-Bridge** — Copyright (c) 2026, NVIDIA CORPORATION, Apache License 2.0. +The hyperparameter values in +[config_nemotron3_nano_30b_a3b_fsdp2.yaml](../../config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml) +and the two Nemotron test configs are adopted from +`src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py` and cross-checked against the model +report. No code was taken from Megatron-Bridge; it served as the configuration reference, and +`src/megatron/bridge/models/nemotronh/nemotron_h_bridge.py` documents the HF parameter mapping that +guided the module naming. + +Files with **no** third-party derivation (written for Modalities, structured after its existing +GPT-2 components): `norms.py`, `nemotron_model_factory.py`, `nemotron_mfu.py`, and the +partitioning algorithm in `nemotron_stages_generator.py`. diff --git a/src/modalities/models/components/mamba2/mamba2_mixer.py b/src/modalities/models/components/mamba2/mamba2_mixer.py index 8fd42aaa7..62f78d7d8 100644 --- a/src/modalities/models/components/mamba2/mamba2_mixer.py +++ b/src/modalities/models/components/mamba2/mamba2_mixer.py @@ -1,3 +1,9 @@ +# Portions of this file are adapted from NVIDIA's Megatron-LM +# (megatron/core/ssm/mamba_mixer.py): the packed [z, x, B, C, dt] input-projection layout, the +# conv1d / A_log / D / dt_bias parameter shapes, and their initialization distributions. +# Copyright (c) 2024, NVIDIA CORPORATION. Copyright (c) 2024, Tri Dao, Albert Gu. +# Licensed under the Apache License, Version 2.0. + """Mamba-2 mixer, the sequence-mixing operator of hybrid Mamba-Transformer models. The layout follows the Megatron-LM reference implementation diff --git a/src/modalities/models/components/mamba2/ssd.py b/src/modalities/models/components/mamba2/ssd.py index 603d2c4ec..c201a71b8 100644 --- a/src/modalities/models/components/mamba2/ssd.py +++ b/src/modalities/models/components/mamba2/ssd.py @@ -1,3 +1,9 @@ +# Portions of this file are adapted from the Mamba-2 reference implementation in +# state-spaces/mamba (https://github.com/state-spaces/mamba): the chunk-parallel SSD block +# decomposition follows `ssd_minimal_discrete` / `segsum`, and GatedRMSNorm reproduces the +# semantics of `mamba_ssm.ops.triton.layernorm_gated.rmsnorm_fn` with norm_before_gate=False. +# Copyright (c) 2024, Tri Dao, Albert Gu. Licensed under the Apache License, Version 2.0. + """Pure-PyTorch Mamba-2 primitives (state space dual / SSD). This module provides a dependency-free, CPU-runnable and ``torch.compile``-friendly diff --git a/src/modalities/models/components/moe/experts.py b/src/modalities/models/components/moe/experts.py index 7ca4fbdd6..03798263a 100644 --- a/src/modalities/models/components/moe/experts.py +++ b/src/modalities/models/components/moe/experts.py @@ -1,3 +1,11 @@ +# Some portions of this implementation are inspired, adapted, or refactored from Meta's +# open-source project TorchTitan (torchtitan/models/common/moe.py::GroupedExperts), +# licensed under the BSD 3-Clause License. Specifically the stacked per-expert weight +# layout and the torch._grouped_mm call over expert-sorted tokens. +# +# The squared ReLU activation follows NVIDIA's Megatron-LM (megatron/core/activations.py), +# Copyright (c) NVIDIA CORPORATION. Licensed under the Apache License, Version 2.0. + """Grouped expert feed-forward networks for mixture-of-experts layers. All experts of a layer are stored in two stacked weight tensors so that the whole layer can be diff --git a/src/modalities/models/components/moe/load_balancing.py b/src/modalities/models/components/moe/load_balancing.py index 109487a6e..80c93fab4 100644 --- a/src/modalities/models/components/moe/load_balancing.py +++ b/src/modalities/models/components/moe/load_balancing.py @@ -1,3 +1,11 @@ +# Portions of this file are adapted from NVIDIA's Megatron-LM +# (megatron/core/transformer/moe/moe_utils.py::get_updated_expert_bias): the sign-based +# per-expert bias update rule. +# Copyright (c) 2025, NVIDIA CORPORATION. Licensed under the Apache License, Version 2.0. +# +# Applying the update as an optimizer step pre-hook follows Meta's open-source project TorchTitan +# (torchtitan/models/common/moe.py), licensed under the BSD 3-Clause License. + """Auxiliary-loss-free load balancing for mixture-of-experts layers. Implements the update rule of "Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts" diff --git a/src/modalities/models/components/moe/moe.py b/src/modalities/models/components/moe/moe.py index 3c4b1ab0a..88ac07bb2 100644 --- a/src/modalities/models/components/moe/moe.py +++ b/src/modalities/models/components/moe/moe.py @@ -1,3 +1,12 @@ +# Portions of this file are adapted from NVIDIA's Megatron-LM +# (megatron/core/transformer/moe/moe_utils.py::switch_load_balancing_loss_func): the +# load-balancing loss formula E * sum_i(f_i * P_i) and its sequence-level variant. +# Copyright (c) 2025, NVIDIA CORPORATION. Licensed under the Apache License, Version 2.0. +# +# The dispatch/combine structure and the expert-bias / token-count buffers are inspired by Meta's +# open-source project TorchTitan (torchtitan/models/common/moe.py::MoE), licensed under the +# BSD 3-Clause License. + """Mixture-of-experts feed-forward layer. The layer combines a :class:`TopKRouter`, a :class:`GroupedExperts` stack and an optional always-on diff --git a/src/modalities/models/components/moe/router.py b/src/modalities/models/components/moe/router.py index 10ad8c037..9902b4d73 100644 --- a/src/modalities/models/components/moe/router.py +++ b/src/modalities/models/components/moe/router.py @@ -1,3 +1,11 @@ +# Portions of this file are adapted from NVIDIA's Megatron-LM +# (megatron/core/transformer/moe/moe_utils.py::topk_routing_with_score_function): the sigmoid +# scoring path, the selection-only expert bias, and the top-k renormalization. +# Copyright (c) 2025, NVIDIA CORPORATION. Licensed under the Apache License, Version 2.0. +# +# The router interface is additionally inspired by Meta's open-source project TorchTitan +# (torchtitan/models/common/moe.py::TokenChoiceTopKRouter), licensed under the BSD 3-Clause License. + """Top-k mixture-of-experts router. Implements the routing scheme used by Nemotron-3 Nano: sigmoid gating with an additive diff --git a/src/modalities/models/nemotron/layer_pattern.py b/src/modalities/models/nemotron/layer_pattern.py index a9f1b543d..9d9dd0a41 100644 --- a/src/modalities/models/nemotron/layer_pattern.py +++ b/src/modalities/models/nemotron/layer_pattern.py @@ -1,3 +1,8 @@ +# The layer-pattern symbols are adopted from NVIDIA's Megatron-LM +# (megatron/core/models/hybrid/hybrid_layer_allocation.py::Symbols) so that pattern strings are +# portable between the two frameworks. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. Licensed under the Apache License, Version 2.0. + """Layer pattern parsing for hybrid Mamba-Transformer models. A hybrid model such as Nemotron-3 Nano is described by a pattern string in which every diff --git a/src/modalities/models/nemotron/nemotron_attention.py b/src/modalities/models/nemotron/nemotron_attention.py index 50518f069..76cf3fdca 100644 --- a/src/modalities/models/nemotron/nemotron_attention.py +++ b/src/modalities/models/nemotron/nemotron_attention.py @@ -1,3 +1,7 @@ +# The grouped-query key/value head repetition follows Meta's Llama implementation +# (https://github.com/facebookresearch/llama), as already adapted in +# modalities/models/gpt2/gpt2_model.py::CausalSelfAttention._repeat_kv. + """Grouped-query causal self-attention with a head dimension decoupled from the model dimension. This cannot reuse the GPT2 attention module: GPT2 derives the head dimension as diff --git a/src/modalities/models/nemotron/nemotron_layer_specs.py b/src/modalities/models/nemotron/nemotron_layer_specs.py index 54fc91124..eb40ae893 100644 --- a/src/modalities/models/nemotron/nemotron_layer_specs.py +++ b/src/modalities/models/nemotron/nemotron_layer_specs.py @@ -1,3 +1,8 @@ +# The layer-spec (declarative builder) pattern is adapted from NVIDIA's Megatron-LM +# (megatron/core/transformer/spec_utils.py::ModuleSpec and +# megatron/core/models/hybrid/hybrid_layer_specs.py). +# Copyright (c) 2023-2026, NVIDIA CORPORATION. Licensed under the Apache License, Version 2.0. + """Layer specifications: configurable builders for the sublayers of a hybrid model. A hybrid model contains dozens of layers of a handful of types. Modalities' component factory diff --git a/src/modalities/models/nemotron/nemotron_layers.py b/src/modalities/models/nemotron/nemotron_layers.py index 76830497a..73f670dc4 100644 --- a/src/modalities/models/nemotron/nemotron_layers.py +++ b/src/modalities/models/nemotron/nemotron_layers.py @@ -1,3 +1,9 @@ +# The single-operator pre-norm residual layer structure is adapted from NVIDIA's Megatron-LM +# (megatron/core/ssm/mamba_layer.py::MambaLayer and +# megatron/core/models/hybrid/hybrid_block.py::HybridStack). +# Copyright (c) 2024-2026, NVIDIA CORPORATION. Copyright (c) 2024, Tri Dao, Albert Gu. +# Licensed under the Apache License, Version 2.0. + """The four residual sublayer types of a hybrid Mamba-Transformer stack. Unlike a classical transformer block, which bundles attention and a feed-forward network, every diff --git a/src/modalities/models/nemotron/nemotron_mlp.py b/src/modalities/models/nemotron/nemotron_mlp.py index 688368ec8..721f09420 100644 --- a/src/modalities/models/nemotron/nemotron_mlp.py +++ b/src/modalities/models/nemotron/nemotron_mlp.py @@ -1,3 +1,6 @@ +# The squared ReLU activation follows NVIDIA's Megatron-LM (megatron/core/activations.py), +# Copyright (c) NVIDIA CORPORATION. Licensed under the Apache License, Version 2.0. + """Non-gated squared-ReLU feed-forward network. Nemotron models use ``ReLU(x)^2`` rather than a gated activation such as SwiGLU. This halves the diff --git a/src/modalities/models/nemotron/nemotron_model.py b/src/modalities/models/nemotron/nemotron_model.py index bb796d8d3..d7947735c 100644 --- a/src/modalities/models/nemotron/nemotron_model.py +++ b/src/modalities/models/nemotron/nemotron_model.py @@ -1,3 +1,11 @@ +# The hybrid model structure (embedding, pattern-driven layer stack, final norm, LM head) is +# adapted from NVIDIA's Megatron-LM (megatron/core/models/hybrid/hybrid_model.py::HybridModel and +# megatron/core/models/hybrid/hybrid_block.py::HybridStack). +# Copyright (c) 2024-2026, NVIDIA CORPORATION. Licensed under the Apache License, Version 2.0. +# +# The module tree deliberately mirrors modalities/models/gpt2/gpt2_model.py::GPT2LLM so that the +# existing FSDP / activation-checkpointing / pipeline components apply unchanged. + """Nemotron-style hybrid Mamba-Transformer language model. The architecture (Nemotron-3 Nano 30B-A3B, arXiv:2512.20848) interleaves Mamba-2 mixers, diff --git a/src/modalities/models/nemotron/nemotron_stages_generator.py b/src/modalities/models/nemotron/nemotron_stages_generator.py index 68cf1d2a9..5e9a75d5c 100644 --- a/src/modalities/models/nemotron/nemotron_stages_generator.py +++ b/src/modalities/models/nemotron/nemotron_stages_generator.py @@ -1,3 +1,7 @@ +# Some portions of this implementation are inspired, adapted, or refactored +# from Meta's open-source project TorchTitan, +# licensed under the BSD 3-Clause License. + """Pipeline stage generation for the Nemotron hybrid model. The GPT2 stages generator assumes every layer costs the same. That is a poor assumption for a hybrid diff --git a/tests/fsdp2_parallelization/nemotron_fsdp2_config.yaml b/tests/fsdp2_parallelization/nemotron_fsdp2_config.yaml index b26521cae..811d6bd66 100644 --- a/tests/fsdp2_parallelization/nemotron_fsdp2_config.yaml +++ b/tests/fsdp2_parallelization/nemotron_fsdp2_config.yaml @@ -1,3 +1,7 @@ +# Hyperparameters are a scaled-down derivative of the Nemotron-3 Nano 30B-A3B architecture +# (arXiv:2512.20848), cross-checked against the Megatron-Bridge recipe +# `src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py`, Copyright (c) 2026, +# NVIDIA CORPORATION, licensed under the Apache License, Version 2.0. # Minimal Nemotron hybrid Mamba-Transformer config for distributed FSDP2 tests. # Exercises all four layer types (Mamba-2, MoE, attention, dense MLP) at a size that fits # comfortably on two GPUs. diff --git a/tests/models/nemotron/test_ssd.py b/tests/models/nemotron/test_ssd.py index f2c2aa56b..a7392db32 100644 --- a/tests/models/nemotron/test_ssd.py +++ b/tests/models/nemotron/test_ssd.py @@ -1,3 +1,7 @@ +# `ssd_recurrent_reference` transcribes the Mamba-2 selective state space recurrence +# (Dao & Gu, 2024, https://arxiv.org/abs/2405.21060) directly from its definition, so that the +# chunk-parallel implementation is validated against the mathematics rather than against a port. + """Correctness tests for the pure-PyTorch Mamba-2 SSD primitives. The chunked scan is the numerical foundation of every Mamba layer, so it is validated against a diff --git a/tests/test_yaml_configs/nemotron_config_initialization.yaml b/tests/test_yaml_configs/nemotron_config_initialization.yaml index 4ff0567e6..617f1c690 100644 --- a/tests/test_yaml_configs/nemotron_config_initialization.yaml +++ b/tests/test_yaml_configs/nemotron_config_initialization.yaml @@ -1,3 +1,7 @@ +# Hyperparameters are a scaled-down derivative of the Nemotron-3 Nano 30B-A3B architecture +# (arXiv:2512.20848), cross-checked against the Megatron-Bridge recipe +# `src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py`, Copyright (c) 2026, +# NVIDIA CORPORATION, licensed under the Apache License, Version 2.0. model: component_key: model variant_key: model_initialized From 2a86c0e568f530c8c70cf82641a3d5886fd0617d Mon Sep 17 00:00:00 2001 From: Max Luebbering <2804731+le1nux@users.noreply.github.com> Date: Tue, 28 Jul 2026 22:18:45 +0200 Subject: [PATCH 3/4] feat(recipe): add a 4-GPU lorem-ipsum config for the Nemotron hybrid model Mirrors config_lorem_ipsum_long_fsdp2.yaml for the Nemotron hybrid Mamba- Transformer. Keeps the full *width* of Nemotron-3 Nano 30B-A3B - dimension 2688, 128 routed experts of dim 1856 with top-6 routing and 2 shared experts, 64 Mamba-2 heads of dim 64, 32 query / 2 KV attention heads of head dim 128 - and cuts only the depth, from 52 to 16 layers, so that it fits on four GPUs. The layer pattern is the first 16 characters of the published pattern, which preserves the Mamba / MoE / attention interleaving and ratio (7 : 7 : 2 here versus 23 : 23 : 6 in the full model). Verified end to end on 4x A100-SXM4-80GB: 162 steps, exit 0, no errors, loss 11.0 -> 0.03, one DCP checkpoint written and reloadable. Measured figures, both variants run: layers pattern parameters active/token peak alloc/GPU 16 MEMEM*EMEMEM*EME 9.67B 0.77B 46.1 GiB 9 MEMEM*EME 5.64B 0.53B 28.9 GiB The 9-layer variant is documented for 40GB cards. Throughput is 5.5 samples/s at 1-2% MFU, which is the expected cost of the native pure-PyTorch state space scan; use ssd_backend=fused for anything beyond a smoke test. Two things the run surfaced, both recorded in the config comments: * A DCP checkpoint of this model is ~109 GiB, so the config checkpoints once near the end. At the reference config's interval of 16 steps this would have written over 1 TB across a 162-step run. * DCPCheckpointSaving._delete_checkpoint uses Path.rmdir(), which only removes empty directories, so rotating a DCP checkpoint always raises OSError("Directory not empty"). That is a pre-existing bug unrelated to this model; the config therefore keeps every checkpoint (k: -1), as the reference config does. Left unfixed deliberately - it deletes directories and deserves its own change plus a test. Also makes NemotronMFUCalculator derive num_active_params from the model when it is omitted. Requiring it as a literal meant every config carried a magic number that silently went stale as soon as the layer pattern or expert count changed. numel() on a DTensor already reports the unsharded size, so the derivation is correct under FSDP2; a list of pipeline stages is rejected with an explicit error, since summing per-stage counts would not give the whole model's active parameters. Co-Authored-By: Claude Opus 5 (1M context) --- ...onfig_lorem_ipsum_nemotron_nano_fsdp2.yaml | 476 ++++++++++++++++++ docs/components/nemotron.md | 13 +- src/modalities/utils/nemotron_mfu.py | 43 +- .../nemotron/test_nemotron_stages_and_mfu.py | 44 ++ 4 files changed, 563 insertions(+), 13 deletions(-) create mode 100644 config_files/training/config_lorem_ipsum_nemotron_nano_fsdp2.yaml diff --git a/config_files/training/config_lorem_ipsum_nemotron_nano_fsdp2.yaml b/config_files/training/config_lorem_ipsum_nemotron_nano_fsdp2.yaml new file mode 100644 index 000000000..4a83e91d0 --- /dev/null +++ b/config_files/training/config_lorem_ipsum_nemotron_nano_fsdp2.yaml @@ -0,0 +1,476 @@ +# Nemotron-3 Nano lorem-ipsum training config for 4x A100 80GB. +# +# This keeps the *full width* of Nemotron-3 Nano 30B-A3B (arXiv:2512.20848) - model dimension 2688, +# 64 Mamba-2 heads of dim 64, 128 routed experts of dim 1856 with top-6 routing and 2 shared +# experts, 32 query / 2 KV attention heads of head dim 128 - and only cuts the depth from 52 to 16 +# layers so that it fits on four GPUs. The layer pattern is the first 16 characters of the published +# pattern, which preserves the Mamba / MoE / attention interleaving and ratio (7 : 7 : 2 here versus +# 23 : 23 : 6 in the full model). +# +# Hyperparameter values are adopted from the Megatron-Bridge recipe +# `src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py`, Copyright (c) 2026, +# NVIDIA CORPORATION, licensed under the Apache License, Version 2.0. +# +# Size and memory, measured on 4x A100-SXM4-80GB (FSDP2 full shard, AdamW, seq 512, micro batch 1): +# +# layers pattern parameters active/token peak allocated / GPU +# 16 MEMEM*EMEMEM*EME 9.67B 0.77B 46.1 GiB +# 9 MEMEM*EME 5.64B 0.53B 28.9 GiB +# +# This file uses the 16-layer variant. For 40GB cards, set `n_layer: 9` and +# `layer_pattern: "MEMEM*EME"`. Every MoE layer costs ~1.30B parameters, a Mamba layer ~39M and an +# attention layer ~23M, so the depth/memory trade-off is driven almost entirely by the `E` count. +# +# Persistent state is 16 bytes per parameter (fp32 shard + fp32 grad + two fp32 Adam moments) +# divided across the 4 ranks; the remainder of the peak is the unsharded bf16 all-gather of the +# largest FSDP unit (one ~1.3B-parameter MoE layer, ~2.4 GiB) plus prefetch and activations. +# +# Run with: +# torchrun --nnodes 1 --nproc_per_node 4 --rdzv-endpoint=0.0.0.0:29555 \ +# $(which modalities) run \ +# --config_file_path config_files/training/config_lorem_ipsum_nemotron_nano_fsdp2.yaml \ +# --experiments_root_path data/experiments +# +# NOTE ON THROUGHPUT: `ssd_backend: native` is the dependency-free pure-PyTorch state space scan. +# It is correct but slow. Install the optional extra (`pip install -e '.[mamba]'`) and switch to +# `ssd_backend: fused` for anything beyond a smoke test. + +settings: + experiment_id: ${modalities_env:experiment_id} + config_file_path: ${modalities_env:config_file_path} + referencing_keys: + sample_key: input_ids + target_key: target_ids + prediction_key: logits + cuda_env: + local_rank: ${cuda_env:LOCAL_RANK} + global_rank: ${cuda_env:RANK} + world_size: ${cuda_env:WORLD_SIZE} + paths: + checkpoint_saving_path: data/checkpoints + train_dataset_path: ./data/lorem_ipsum_long.pbin + test_dataset_path: ./data/lorem_ipsum.pbin + experiments_root_path: ${modalities_env:experiments_root_path} + intervals: + training_log_interval_in_steps: 1 + # A DCP checkpoint of this model is ~109 GiB (9.67B parameters plus fp32 Adam state, written by + # all 4 ranks), so checkpoint once near the end rather than every 16 steps - the latter would + # write over 1 TB across a 162-step run. Raise the frequency only if you have the disk for it. + checkpointing_interval_in_steps: 160 + evaluation_interval_in_steps: 16 + consistency_enforcement: + enforce_tokens_per_step_consistency: false + enforce_last_step_logged: false + enforce_last_step_evaluated: false + enforce_last_step_checkpointed: false + step_profile: + gradient_accumulation_steps: 1 + local_train_micro_batch_size: 1 + # 512 with a Mamba chunk size of 128 gives four scan chunks, so the inter-chunk state passing + # is actually exercised rather than degenerating to a single chunk. + sequence_length: 512 + dp_degree: + instance_key: dp_degree + pass_type: BY_REFERENCE + training_target: + num_target_tokens: + component_key: number_conversion + variant_key: num_tokens_from_packed_mem_map_dataset_continuous + config: + dataset_path: ${settings.paths.train_dataset_path} + sequence_length: ${settings.step_profile.sequence_length} + dp_degree: + instance_key: dp_degree + pass_type: BY_REFERENCE + local_micro_batch_size: ${settings.step_profile.local_train_micro_batch_size} + gradient_accumulation_steps: ${settings.step_profile.gradient_accumulation_steps} + num_target_steps: + component_key: number_conversion + variant_key: num_steps_from_num_tokens + config: + dp_degree: + instance_key: dp_degree + pass_type: BY_REFERENCE + local_micro_batch_size: ${settings.step_profile.local_train_micro_batch_size} + global_num_tokens: ${settings.training_target.num_target_tokens} + sequence_length: ${settings.step_profile.sequence_length} + gradient_accumulation_steps: ${settings.step_profile.gradient_accumulation_steps} + training_progress: + global_num_seen_tokens: 0 + num_seen_steps: 0 + num_seen_samples: 0 + last_step: -1 + +collate_fn: + component_key: collate_fn + variant_key: gpt_2_llm_collator + config: + sample_key: ${settings.referencing_keys.sample_key} + target_key: ${settings.referencing_keys.target_key} + +train_dataset: + component_key: dataset + variant_key: packed_mem_map_dataset_continuous + config: + raw_data_path: ${settings.paths.train_dataset_path} + sequence_length: ${settings.step_profile.sequence_length} + sample_key: ${settings.referencing_keys.sample_key} + +train_dataloader: + component_key: data_loader + variant_key: default + config: + num_workers: 2 + pin_memory: true + dataloader_tag: train + dataset: + instance_key: train_dataset + pass_type: BY_REFERENCE + batch_sampler: + component_key: batch_sampler + variant_key: default + config: + batch_size: ${settings.step_profile.local_train_micro_batch_size} + drop_last: true + sampler: + component_key: sampler + variant_key: resumable_distributed_sampler + config: + dataset: + instance_key: train_dataset + pass_type: BY_REFERENCE + rank: ${settings.cuda_env.global_rank} + num_replicas: ${settings.cuda_env.world_size} + shuffle: true + seed: 42 + drop_last: true + skip_num_global_samples: ${settings.training_progress.num_seen_samples} + collate_fn: + instance_key: collate_fn + pass_type: BY_REFERENCE + +test_dataset: + component_key: dataset + variant_key: packed_mem_map_dataset_continuous + config: + raw_data_path: ${settings.paths.test_dataset_path} + sequence_length: ${settings.step_profile.sequence_length} + sample_key: ${settings.referencing_keys.sample_key} + +test_dataloader: + component_key: data_loader + variant_key: default + config: + num_workers: 2 + pin_memory: true + dataloader_tag: test + dataset: + instance_key: test_dataset + pass_type: BY_REFERENCE + batch_sampler: + component_key: batch_sampler + variant_key: default + config: + batch_size: ${settings.step_profile.local_train_micro_batch_size} + drop_last: true + sampler: + component_key: sampler + variant_key: distributed_sampler + config: + rank: ${settings.cuda_env.global_rank} + num_replicas: ${settings.cuda_env.world_size} + shuffle: false + drop_last: true + dataset: + instance_key: test_dataset + pass_type: BY_REFERENCE + collate_fn: + instance_key: collate_fn + pass_type: BY_REFERENCE + +eval_dataloaders: + - instance_key: test_dataloader + pass_type: BY_REFERENCE + +checkpoint_saving: + component_key: checkpoint_saving + variant_key: default + config: + checkpoint_saving_strategy: + component_key: checkpoint_saving_strategy + variant_key: save_k_most_recent_checkpoints_strategy + config: + # -1 keeps every checkpoint, as in config_lorem_ipsum_long_fsdp2.yaml. Do not set this to a + # positive value with the `dcp` execution: DCPCheckpointSaving._delete_checkpoint uses + # Path.rmdir(), which only removes empty directories, so rotating a DCP checkpoint raises + # OSError("Directory not empty"). That is a pre-existing bug, unrelated to this model. + k: -1 + checkpoint_saving_execution: + component_key: checkpoint_saving_execution + variant_key: dcp + config: + checkpoint_path: ${settings.paths.checkpoint_saving_path} + global_rank: ${settings.cuda_env.global_rank} + experiment_id: ${settings.experiment_id} + +# Language modelling loss plus the MoE load-balancing penalty. The penalty is computed inside each +# MoE layer (coefficient 1e-4, per the model report) and summed by the model into `moe_aux_loss`. +# The *primary* balancing mechanism is the auxiliary-loss-free expert bias; see the optimizer. +loss_fn: + component_key: loss + variant_key: weighted_sum + config: + weights: [1.0, 1.0] + losses: + - component_key: loss + variant_key: clm_cross_entropy_loss + config: + target_key: ${settings.referencing_keys.target_key} + prediction_key: ${settings.referencing_keys.prediction_key} + - component_key: loss + variant_key: moe_aux_loss + config: + prediction_key: moe_aux_loss + +device_mesh: + component_key: device_mesh + variant_key: default + config: + device_type: cuda + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + world_size: ${settings.cuda_env.world_size} + +dp_degree: + component_key: number_conversion + variant_key: parallel_degree + config: + device_mesh: + instance_key: device_mesh + pass_type: BY_REFERENCE + parallelism_methods: [dp_shard, dp_replicate] + +app_state: + component_key: app_state + variant_key: raw + config: + model: + instance_key: initialized_model + pass_type: BY_REFERENCE + optimizer: + instance_key: optimizer + pass_type: BY_REFERENCE + lr_scheduler: + instance_key: lr_scheduler + pass_type: BY_REFERENCE + +initialized_model: + component_key: model + variant_key: model_initialized + config: + model: + instance_key: fsdp_model + pass_type: BY_REFERENCE + model_initializer: + component_key: model_initialization + variant_key: composed + config: + model_type: nemotron + weight_init_type: scaled + mean: 0.0 + # "auto" resolves to sqrt(2 / (5 * n_embd)) = 0.0172, which is what the reference recipe + # uses as init_method_std for this width. + std: auto + hidden_dim: ${model_raw.config.n_embd} + num_layers: ${model_raw.config.n_layer} + +fsdp_model: + component_key: model + variant_key: fsdp2_wrapped + config: + model: + instance_key: activation_checkpointed_model + pass_type: BY_REFERENCE + device_mesh: + instance_key: device_mesh + pass_type: BY_REFERENCE + mixed_precision_settings: + param_dtype: BF_16 + reduce_dtype: FP_32 + # One FSDP unit per layer, so only a single ~1.3B-parameter MoE block is unsharded at a time. + layers_per_fsdp_unit: 1 + block_names: [Mamba2Layer, NemotronMoELayer, NemotronAttentionLayer, NemotronMLPLayer] + +activation_checkpointed_model: + component_key: model + variant_key: activation_checkpointed + config: + model: + instance_key: model_raw + pass_type: BY_REFERENCE + ac_variant: full_activation_checkpointing + layers_fqn: transformer.h + ac_fun_params: {} + +model_raw: + component_key: model + variant_key: nemotron + config: + use_meta_device: true + sample_key: ${settings.referencing_keys.sample_key} + prediction_key: ${settings.referencing_keys.prediction_key} + aux_loss_key: moe_aux_loss + sequence_length: ${settings.step_profile.sequence_length} + # GPT-2 vocab of 50257, padded to a multiple of 128 (the lorem ipsum data is GPT-2 tokenized). + vocab_size: 50304 + n_embd: 2688 + # Full Nemotron-3 Nano is 52 layers; trimmed to 16 to fit four GPUs. Keep this in sync with + # the length of layer_pattern - the model validates it. + n_layer: 16 + layer_pattern: "MEMEM*EMEMEM*EME" + use_weight_tying: false + lm_head_norm_config: &nemotron_norm_config + norm_type: pytorch_rms_norm + config: + normalized_shape: ${model_raw.config.n_embd} + eps: 1e-5 + layer_specs: + "M": + component_key: nemotron_layer_spec + variant_key: mamba2 + config: + n_embd: ${model_raw.config.n_embd} + mamba_n_heads: 64 + mamba_head_dim: 64 + mamba_state_dim: 128 + mamba_n_groups: 8 + d_conv: 4 + chunk_size: 128 + ssd_backend: native # switch to `fused` once mamba-ssm / causal-conv1d are installed + norm_config: *nemotron_norm_config + "E": + component_key: nemotron_layer_spec + variant_key: moe + config: + n_embd: ${model_raw.config.n_embd} + num_experts: 128 + moe_ffn_hidden: 1856 + top_k: 6 + route_scale: 2.5 + score_function: sigmoid + use_expert_bias: true + router_dtype: float32 + num_shared_experts: 2 + shared_expert_ffn_hidden_per_expert: 1856 # -> one fused MLP of 3712 hidden units + aux_loss_coeff: 1.0e-4 + experts_backend: grouped_mm + norm_config: *nemotron_norm_config + "*": + component_key: nemotron_layer_spec + variant_key: attention + config: + n_embd: ${model_raw.config.n_embd} + n_head_q: 32 + n_head_kv: 2 + head_dim: 128 + attention_implementation: pytorch_flash + norm_config: *nemotron_norm_config + +lr_scheduler: + component_key: scheduler + variant_key: onecycle_lr + config: + optimizer: + instance_key: optimizer + pass_type: BY_REFERENCE + max_lr: 4.5e-4 + div_factor: 10 + final_div_factor: 1 + total_steps: ${settings.training_target.num_target_steps} + pct_start: 0.1 + anneal_strategy: cos + last_epoch: ${settings.training_progress.last_step} + +# The `moe_load_balanced` decorator adds the auxiliary-loss-free expert bias update as an optimizer +# step pre-hook. Pinning it to the optimizer step (rather than the forward pass) is what makes it +# correct under gradient accumulation: token counts accumulate over micro-batches and are reduced +# across data-parallel ranks exactly once per step. +optimizer: + component_key: optimizer + variant_key: moe_load_balanced + config: + expert_bias_update_rate: 1.0e-3 + model: + instance_key: initialized_model + pass_type: BY_REFERENCE + device_mesh: + instance_key: device_mesh + pass_type: BY_REFERENCE + optimizer: + component_key: optimizer + variant_key: adam_w + config: + lr: 4.5e-4 + betas: [0.9, 0.95] + eps: 1e-8 + weight_decay: 0.1 + # The state space parameters (A_log, D, dt_bias, conv1d) parameterize the SSM dynamics and + # the router gate decides expert assignment; decaying either destabilizes training. + weight_decay_groups_excluded: [embedding, layernorm, ssm, router] + wrapped_model: + instance_key: initialized_model + pass_type: BY_REFERENCE + +gradient_clipper: + component_key: gradient_clipper + variant_key: fsdp2 + config: + wrapped_model: + instance_key: initialized_model + pass_type: BY_REFERENCE + norm_type: P2_NORM + max_norm: 1.0 + device_mesh: + instance_key: device_mesh + pass_type: BY_REFERENCE + +progress_subscriber: + component_key: progress_subscriber + variant_key: rich + config: + global_rank: ${settings.cuda_env.global_rank} + num_seen_steps: ${settings.training_progress.num_seen_steps} + num_target_steps: ${settings.training_target.num_target_steps} + train_dataloader_tag: ${train_dataloader.config.dataloader_tag} + eval_dataloaders: + instance_key: eval_dataloaders + pass_type: BY_REFERENCE + +evaluation_subscriber: + component_key: results_subscriber + variant_key: wandb + config: + global_rank: ${settings.cuda_env.global_rank} + project: modalities_nemotron_tests + mode: OFFLINE + experiment_id: ${settings.experiment_id} + directory: wandb_storage + config_file_path: ${settings.config_file_path} + +# num_active_params is derived from the model, so it stays correct if the layer pattern changes. +mfu_calculator: + component_key: mfu_calculator + variant_key: nemotron + config: + layer_pattern: ${model_raw.config.layer_pattern} + sequence_length: ${settings.step_profile.sequence_length} + n_embd: ${model_raw.config.n_embd} + n_head_q: 32 + head_dim: 128 + world_size: ${settings.cuda_env.world_size} + model_parts: + instance_key: initialized_model + pass_type: BY_REFERENCE + device_mesh: + instance_key: device_mesh + pass_type: BY_REFERENCE diff --git a/docs/components/nemotron.md b/docs/components/nemotron.md index d114bb13b..9cef4ee40 100644 --- a/docs/components/nemotron.md +++ b/docs/components/nemotron.md @@ -4,8 +4,17 @@ Modalities supports hybrid Mamba-Transformer architectures with sparse mixture-o layers, as introduced by NVIDIA's Nemotron-H and Nemotron-3 Nano families. The reference target is **Nemotron-3 Nano 30B-A3B** ([arXiv:2512.20848](https://arxiv.org/abs/2512.20848)). -A ready-to-run pretraining configuration is at -[config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml](../../config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml). +Two ready-to-run configurations are provided: + +| Config | Scale | Verified on | +|--------|-------|-------------| +| [config_nemotron3_nano_30b_a3b_fsdp2.yaml](../../config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml) | Full 52-layer 30B-A3B | Not run (needs tensor/expert parallelism for real throughput) | +| [config_lorem_ipsum_nemotron_nano_fsdp2.yaml](../../config_files/training/config_lorem_ipsum_nemotron_nano_fsdp2.yaml) | Full width, 16 layers, 9.67B params | 4x A100-SXM4-80GB, 162 steps, 46.1 GiB peak per GPU | + +The lorem-ipsum config keeps every width hyperparameter of the real model (dimension 2688, 128 +experts, 64 Mamba heads, 32/2 attention heads) and only trims the depth, so it exercises the real +architecture end to end. Shorten `layer_pattern` to `"MEMEM*EME"` with `n_layer: 9` for 40GB cards +(5.64B parameters, 28.9 GiB peak per GPU). ## Architecture diff --git a/src/modalities/utils/nemotron_mfu.py b/src/modalities/utils/nemotron_mfu.py index 8d8bf9bfc..4b8461f9a 100644 --- a/src/modalities/utils/nemotron_mfu.py +++ b/src/modalities/utils/nemotron_mfu.py @@ -34,8 +34,9 @@ class NemotronMFUCalculatorConfig(BaseModel): n_embd (int): Model dimension. n_head_q (int): Number of query heads of the attention layers. head_dim (int): Head dimension of the attention layers. - num_active_params (int): Number of parameters visited per token. See - :meth:`NemotronMFUCalculator.count_active_parameters`. + num_active_params (int | None): Number of parameters visited per token. Derived from + ``model_parts`` when omitted, which is what you normally want: hardcoding it in a config + silently goes stale as soon as the layer pattern or expert count changes. world_size (int): Number of ranks. model_parts (nn.Module | list[nn.Module]): The wrapped model (or pipeline stages). device_mesh (DeviceMesh | None): The device mesh, if any. @@ -46,7 +47,7 @@ class NemotronMFUCalculatorConfig(BaseModel): n_embd: Annotated[int, Field(strict=True, ge=1)] n_head_q: Annotated[int, Field(strict=True, ge=1)] head_dim: Annotated[int, Field(strict=True, ge=1)] - num_active_params: Annotated[int, Field(strict=True, ge=1)] + num_active_params: Optional[Annotated[int, Field(strict=True, ge=1)]] = None world_size: Annotated[int, Field(strict=True, ge=1)] model_parts: PydanticPytorchModuleOrListType device_mesh: Optional[object] = None @@ -65,9 +66,9 @@ def __init__( n_embd: int, n_head_q: int, head_dim: int, - num_active_params: int, - world_size: int, - model_parts: torch.nn.Module | list[torch.nn.Module], + num_active_params: Optional[int] = None, + world_size: int = 1, + model_parts: torch.nn.Module | list[torch.nn.Module] = None, device_mesh: Optional[object] = None, ): """ @@ -79,12 +80,17 @@ def __init__( n_embd (int): Model dimension. n_head_q (int): Number of query heads of the attention layers. head_dim (int): Head dimension of the attention layers. - num_active_params (int): Number of parameters visited per token. + num_active_params (int | None): Number of parameters visited per token. Derived from + ``model_parts`` when omitted. world_size (int): Number of ranks. model_parts (nn.Module | list[nn.Module]): The wrapped model or pipeline stages. device_mesh (DeviceMesh | None): The device mesh, if any. + + Raises: + ValueError: If ``num_active_params`` is omitted and cannot be derived. """ - del device_mesh # only needed by the parameter-counting calculators + if num_active_params is None: + num_active_params = NemotronMFUCalculator.count_active_parameters(model_parts) self._sequence_length = sequence_length self._num_attention_layers = count_layers_by_type(layer_pattern)[LayerSymbol.ATTENTION] self._theoretical_flops_per_token = NemotronMFUCalculator._get_theoretical_flops_per_token( @@ -97,7 +103,7 @@ def __init__( self._theoretical_gpu_peak_performance = MFUCalculatorABC._get_theoretical_gpu_peak_performance( model_parts, world_size ) - del n_embd # part of the public config for symmetry with the GPT2 calculator + del n_embd, device_mesh # part of the public config for symmetry with the GPT2 calculator @staticmethod def _get_theoretical_flops_per_token( @@ -138,15 +144,30 @@ def count_active_parameters(model: torch.nn.Module) -> int: ``num_experts`` are evaluated per token, so their contribution is scaled accordingly. The router, the shared experts and all dense layers are always active. + Works on plain and FSDP2-wrapped models: ``numel()`` on a ``DTensor`` reports the unsharded + size, so the result is the global count regardless of sharding. + Args: - model (nn.Module): The (unwrapped) model. + model (nn.Module | list[nn.Module]): The model. A list of pipeline stages is rejected, + because summing per-stage counts would not give the whole model's active parameters. + + Raises: + ValueError: If ``model`` is None or a list of pipeline stages. Returns: int: The number of active parameters. """ from modalities.models.components.moe.moe import MoE - total = sum(p.numel() for p in model.parameters()) + if model is None: + raise ValueError("num_active_params was omitted but no model was provided to derive it from.") + if isinstance(model, list): + raise ValueError( + "Cannot derive num_active_params from a list of pipeline stages, since each stage holds " + "only a subset of the layers. Pass num_active_params explicitly." + ) + + total = sum(p.numel() for p in model.parameters() if p.requires_grad) for module in model.modules(): if not isinstance(module, MoE): continue diff --git a/tests/models/nemotron/test_nemotron_stages_and_mfu.py b/tests/models/nemotron/test_nemotron_stages_and_mfu.py index eed66b2dc..cf239ace3 100644 --- a/tests/models/nemotron/test_nemotron_stages_and_mfu.py +++ b/tests/models/nemotron/test_nemotron_stages_and_mfu.py @@ -346,3 +346,47 @@ def test_full_config_layer_specs_validate(): attention_spec = NemotronAttentionLayerSpec(**{**config["layer_specs"]["*"]["config"], "norm_config": norm_config}) assert attention_spec.config.n_head_q // attention_spec.config.n_head_kv == 16 + + +def test_mfu_calculator_derives_active_params_when_omitted(monkeypatch): + # Hardcoding num_active_params in a config silently goes stale when the layer pattern or the + # expert count changes, so omitting it must derive the value from the model. The GPU peak + # performance lookup is stubbed out: it needs an FSDP-wrapped model, which is unrelated to the + # derivation being tested here (the wrapped path is covered by the 4-GPU config run). + monkeypatch.setattr( + "modalities.utils.mfu.MFUCalculatorABC._get_theoretical_gpu_peak_performance", + staticmethod(lambda model_parts, world_size: 1.0), + ) + model = _make_model(layer_pattern="ME*-", n_layer=4) + expected = NemotronMFUCalculator.count_active_parameters(model) + + common = dict(layer_pattern="ME*-", sequence_length=32, n_embd=model.n_embd, n_head_q=4, head_dim=16) + explicit = NemotronMFUCalculator(**common, num_active_params=expected, world_size=1, model_parts=model) + derived = NemotronMFUCalculator(**common, world_size=1, model_parts=model) + assert derived._theoretical_flops_per_token == explicit._theoretical_flops_per_token + # And the derived value really is the sparse one, not the total parameter count. + assert expected < sum(p.numel() for p in model.parameters()) + + +def test_mfu_calculator_rejects_deriving_active_params_from_pipeline_stages(): + # Each stage holds only a subset of the layers, so summing per-stage counts would be wrong. + model = _make_model(layer_pattern="ME", n_layer=2) + with pytest.raises(ValueError, match="list of pipeline stages"): + NemotronMFUCalculator.count_active_parameters([model, model]) + with pytest.raises(ValueError, match="no model was provided"): + NemotronMFUCalculator.count_active_parameters(None) + + +def test_mfu_calculator_config_allows_omitting_active_params(): + from modalities.utils.nemotron_mfu import NemotronMFUCalculatorConfig + + config = NemotronMFUCalculatorConfig( + layer_pattern="ME", + sequence_length=32, + n_embd=128, + n_head_q=4, + head_dim=16, + world_size=1, + model_parts=_make_model(layer_pattern="ME", n_layer=2), + ) + assert config.num_active_params is None From 90ceb8d0dc76e1c5de678b27c1745d6455aa705e Mon Sep 17 00:00:00 2001 From: Max Luebbering <2804731+le1nux@users.noreply.github.com> Date: Wed, 29 Jul 2026 12:00:26 +0200 Subject: [PATCH 4/4] fix(ckpt): delete DCP checkpoints recursively so rotation works DCPCheckpointSaving._delete_checkpoint removed checkpoints with Path.rmdir(), which only removes *empty* directories. A distributed checkpoint is a directory of per-rank shard files (.metadata plus one ___0.distcp per rank), so deleting one always raised OSError("Directory not empty"). The effect was that every rotating checkpoint strategy was broken whenever it was combined with the dcp execution: save_k_most_recent_checkpoints_strategy with k >= 1 and keep_every_k_steps_and_m_most_recent_checkpointing_strategy both crash training at the first rotation, once enough checkpoints exist for one to be evicted. Only k = -1 (keep everything) worked, which is why the existing lorem-ipsum configs never hit it. Deletion is now recursive. Since that turns a no-op into a real recursive delete, it is guarded: the target must exist, must be a directory, and must resolve to a path located directly inside the configured checkpoint_path. The path itself is not user input - it is rebuilt from CHECKPOINT_FOLDER_STRUCTURE for a TrainingProgress this instance previously saved - but the guard makes that invariant explicit rather than implicit, and conservatively refuses a checkpoint folder that symlinks somewhere else. Adds tests/checkpointing/test_dcp_checkpoint_deletion.py, which builds folders shaped like real DCP checkpoints (shard files plus a nested directory). Five of its seven tests fail against the old rmdir implementation, including an end-to-end rotation through SaveKMostRecentCheckpointsStrategy. The FSDP1 execution is untouched: it deletes individual files with unlink(), which is correct for its single-file checkpoints. Also relaxes the Mamba fused-backend error assertion, which named both optional packages and so broke when only one of mamba-ssm / causal-conv1d was installed, and refreshes the checkpointing comment in the Nemotron lorem-ipsum config now that rotation is no longer a trap. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG_DEV.md | 10 +- ...onfig_lorem_ipsum_nemotron_nano_fsdp2.yaml | 8 +- .../fsdp/fsdp_checkpoint_saving.py | 26 ++- .../test_dcp_checkpoint_deletion.py | 191 ++++++++++++++++++ tests/models/nemotron/test_mamba2_mixer.py | 6 +- 5 files changed, 230 insertions(+), 11 deletions(-) create mode 100644 tests/checkpointing/test_dcp_checkpoint_deletion.py diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 3cab52010..412de3939 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -253,7 +253,15 @@ See [docs/components/nemotron.md](docs/components/nemotron.md) for the full guid `Mamba2Mixer.reset_parameters` instead, matching the reference distributions. * Ready-to-run config: `config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml`. -**Bug fixes (pre-existing, surfaced by the distributed tests for this feature):** +**Bug fixes (pre-existing, surfaced while building and running this feature):** +* `DCPCheckpointSaving._delete_checkpoint` removed checkpoints with `Path.rmdir()`, which only works + on *empty* directories. A distributed checkpoint is a directory of per-rank shard files, so + deleting one always raised `OSError("Directory not empty")`. Every rotating checkpoint strategy + (`save_k_most_recent_checkpoints_strategy` with `k >= 1`, + `keep_every_k_steps_and_m_most_recent_checkpointing_strategy`) was therefore broken when combined + with the `dcp` execution: training crashed at the first rotation. Deletion is now recursive, and + guarded so that it can only ever remove a directory located directly inside the configured + checkpoint path. * `NamedParameterwiseNormalInitialization` only stripped torch.compile's `_orig_mod.` prefix from parameter names before matching the initialization filters. Activation checkpointing and FSDP1 insert their own segments (`_checkpoint_wrapped_module.`, `_fsdp_wrapped_module.`), so any config diff --git a/config_files/training/config_lorem_ipsum_nemotron_nano_fsdp2.yaml b/config_files/training/config_lorem_ipsum_nemotron_nano_fsdp2.yaml index 4a83e91d0..d571cdd74 100644 --- a/config_files/training/config_lorem_ipsum_nemotron_nano_fsdp2.yaml +++ b/config_files/training/config_lorem_ipsum_nemotron_nano_fsdp2.yaml @@ -200,10 +200,10 @@ checkpoint_saving: component_key: checkpoint_saving_strategy variant_key: save_k_most_recent_checkpoints_strategy config: - # -1 keeps every checkpoint, as in config_lorem_ipsum_long_fsdp2.yaml. Do not set this to a - # positive value with the `dcp` execution: DCPCheckpointSaving._delete_checkpoint uses - # Path.rmdir(), which only removes empty directories, so rotating a DCP checkpoint raises - # OSError("Directory not empty"). That is a pre-existing bug, unrelated to this model. + # -1 keeps every checkpoint, as in config_lorem_ipsum_long_fsdp2.yaml. Since only one + # checkpoint is written here (see checkpointing_interval_in_steps), rotation never kicks in. + # Setting a positive k does work - checkpoint rotation with the `dcp` execution is fixed - + # but each checkpoint of this model is ~109 GiB, so prefer a coarse interval over rotation. k: -1 checkpoint_saving_execution: component_key: checkpoint_saving_execution diff --git a/src/modalities/checkpointing/fsdp/fsdp_checkpoint_saving.py b/src/modalities/checkpointing/fsdp/fsdp_checkpoint_saving.py index 0caeeb52c..e5b8ad9e0 100644 --- a/src/modalities/checkpointing/fsdp/fsdp_checkpoint_saving.py +++ b/src/modalities/checkpointing/fsdp/fsdp_checkpoint_saving.py @@ -1,4 +1,5 @@ import json +import shutil from enum import Enum from pathlib import Path @@ -273,10 +274,27 @@ def _delete_checkpoint(self, training_progress: TrainingProgress): num_target_steps=training_progress.num_target_steps, num_target_tokens=training_progress.num_target_tokens, ) - if folder_path_to_delete.exists(): - # unlink removes the file - folder_path_to_delete.rmdir() - else: + if not folder_path_to_delete.exists(): raise CheckpointingError( f"Checkpoint folder {folder_path_to_delete} could not be removed. It does not exist!" ) + if not folder_path_to_delete.is_dir(): + raise CheckpointingError(f"Checkpoint path {folder_path_to_delete} is not a directory and was not removed.") + # A distributed checkpoint is a *directory* of per-rank shard files, so it has to be removed + # recursively. Path.rmdir() only removes empty directories and therefore always raised + # OSError("Directory not empty") here, which broke every rotating checkpoint strategy + # (save_k_most_recent_checkpoints_strategy with k >= 1, keep_every_k_steps_...) when combined + # with this DCP execution. + # + # The path is not taken from user input: it is rebuilt from CHECKPOINT_FOLDER_STRUCTURE for a + # TrainingProgress that this instance previously saved. The guard below makes that explicit + # rather than implicit, so a recursive delete can never escape the configured checkpoint + # directory. + expected_parent = self.checkpoint_path.resolve() + resolved_path = folder_path_to_delete.resolve() + if resolved_path.parent != expected_parent: + raise CheckpointingError( + f"Refusing to delete {resolved_path}: it is not directly inside the configured " + f"checkpoint path {expected_parent}." + ) + shutil.rmtree(resolved_path) diff --git a/tests/checkpointing/test_dcp_checkpoint_deletion.py b/tests/checkpointing/test_dcp_checkpoint_deletion.py new file mode 100644 index 000000000..7529c9cc4 --- /dev/null +++ b/tests/checkpointing/test_dcp_checkpoint_deletion.py @@ -0,0 +1,191 @@ +"""Tests for DCP checkpoint deletion. + +A distributed checkpoint is a *directory* of per-rank shard files. Deleting one therefore has to be +recursive; ``Path.rmdir()`` only removes empty directories and always raised +``OSError("Directory not empty")``, which silently broke every rotating checkpoint strategy that was +combined with the DCP execution. +""" + +import json +from pathlib import Path + +import pytest + +from modalities.checkpointing.fsdp.fsdp_checkpoint_saving import DCPCheckpointSaving +from modalities.exceptions import CheckpointingError +from modalities.training.training_progress import TrainingProgress + +EXPERIMENT_ID = "test_experiment" + + +def _training_progress() -> TrainingProgress: + return TrainingProgress( + num_seen_steps_current_run=8, + num_seen_tokens_current_run=64, + num_target_steps=16, + num_target_tokens=128, + ) + + +def _make_saving(checkpoint_path: Path, global_rank: int = 0) -> DCPCheckpointSaving: + return DCPCheckpointSaving( + checkpoint_path=checkpoint_path, + experiment_id=EXPERIMENT_ID, + global_rank=global_rank, + ) + + +def _create_checkpoint_folder(saving: DCPCheckpointSaving, training_progress: TrainingProgress) -> Path: + """Creates a folder shaped like a real DCP checkpoint: a directory containing shard files.""" + folder = saving._get_checkpointing_folder_path( + experiment_id=EXPERIMENT_ID, + num_seen_steps=training_progress.num_seen_steps_total, + num_seen_tokens=training_progress.num_seen_tokens_total, + num_target_steps=training_progress.num_target_steps, + num_target_tokens=training_progress.num_target_tokens, + ) + folder.mkdir(parents=True) + # This is what makes rmdir() fail: DCP writes one shard file per rank plus metadata. + (folder / ".metadata").write_bytes(b"metadata") + for rank in range(4): + (folder / f"__{rank}_0.distcp").write_bytes(b"shard") + (folder / "nested").mkdir() + (folder / "nested" / "extra.bin").write_bytes(b"nested") + return folder + + +def test_deletes_a_non_empty_checkpoint_directory(tmp_path): + # The regression: a DCP checkpoint folder is never empty, so deletion must be recursive. + saving = _make_saving(tmp_path) + training_progress = _training_progress() + folder = _create_checkpoint_folder(saving, training_progress) + assert any(folder.iterdir()), "test fixture must produce a non-empty directory" + + saving._delete_checkpoint(training_progress=training_progress) + + assert not folder.exists() + # Only the checkpoint folder goes away; the surrounding checkpoint directory stays. + assert tmp_path.exists() + + +def test_deletion_leaves_sibling_checkpoints_untouched(tmp_path): + saving = _make_saving(tmp_path) + to_delete = _training_progress() + keep = TrainingProgress( + num_seen_steps_current_run=16, + num_seen_tokens_current_run=128, + num_target_steps=16, + num_target_tokens=128, + ) + folder_to_delete = _create_checkpoint_folder(saving, to_delete) + folder_to_keep = _create_checkpoint_folder(saving, keep) + info_file = tmp_path / "last_checkpoint_info.json" + info_file.write_text(json.dumps({"checkpoint_folder_path": str(folder_to_keep)})) + + saving._delete_checkpoint(training_progress=to_delete) + + assert not folder_to_delete.exists() + assert folder_to_keep.exists() + assert sorted(p.name for p in folder_to_keep.iterdir()) == [ + ".metadata", + "__0_0.distcp", + "__1_0.distcp", + "__2_0.distcp", + "__3_0.distcp", + "nested", + ] + assert info_file.exists() + + +def test_non_zero_rank_does_not_delete(tmp_path): + # Only rank 0 removes checkpoints; every other rank must be a no-op. + saving = _make_saving(tmp_path, global_rank=1) + training_progress = _training_progress() + folder = _create_checkpoint_folder(saving, training_progress) + + saving._delete_checkpoint(training_progress=training_progress) + + assert folder.exists() + + +def test_missing_checkpoint_raises(tmp_path): + saving = _make_saving(tmp_path) + with pytest.raises(CheckpointingError, match="does not exist"): + saving._delete_checkpoint(training_progress=_training_progress()) + + +def test_refuses_to_delete_a_file(tmp_path): + saving = _make_saving(tmp_path) + training_progress = _training_progress() + folder = saving._get_checkpointing_folder_path( + experiment_id=EXPERIMENT_ID, + num_seen_steps=training_progress.num_seen_steps_total, + num_seen_tokens=training_progress.num_seen_tokens_total, + num_target_steps=training_progress.num_target_steps, + num_target_tokens=training_progress.num_target_tokens, + ) + folder.parent.mkdir(parents=True, exist_ok=True) + folder.write_bytes(b"not a directory") + + with pytest.raises(CheckpointingError, match="is not a directory"): + saving._delete_checkpoint(training_progress=training_progress) + assert folder.exists() + + +def test_refuses_to_delete_outside_the_configured_checkpoint_path(tmp_path): + # A recursive delete must not be able to escape the configured checkpoint directory, even if the + # experiment id were ever to contain path separators. + outside = tmp_path / "outside" + outside.mkdir() + (outside / "precious.bin").write_bytes(b"do not delete") + + checkpoint_path = tmp_path / "checkpoints" + checkpoint_path.mkdir() + saving = DCPCheckpointSaving( + checkpoint_path=checkpoint_path, + experiment_id="../outside/escaped", + global_rank=0, + ) + training_progress = _training_progress() + folder = saving._get_checkpointing_folder_path( + experiment_id="../outside/escaped", + num_seen_steps=training_progress.num_seen_steps_total, + num_seen_tokens=training_progress.num_seen_tokens_total, + num_target_steps=training_progress.num_target_steps, + num_target_tokens=training_progress.num_target_tokens, + ) + folder.mkdir(parents=True) + (folder / "shard.distcp").write_bytes(b"shard") + + with pytest.raises(CheckpointingError, match="Refusing to delete"): + saving._delete_checkpoint(training_progress=training_progress) + assert (outside / "precious.bin").exists() + assert folder.exists() + + +def test_rotation_end_to_end_with_the_save_k_most_recent_strategy(tmp_path): + # The failure this fixes only showed up through the strategy: it asks the execution to delete the + # checkpoint that just fell out of the k-most-recent window. + from modalities.checkpointing.checkpoint_saving_strategies import SaveKMostRecentCheckpointsStrategy + + saving = _make_saving(tmp_path) + strategy = SaveKMostRecentCheckpointsStrategy(k=1) + + first = _training_progress() + second = TrainingProgress( + num_seen_steps_current_run=16, + num_seen_tokens_current_run=128, + num_target_steps=16, + num_target_tokens=128, + ) + first_folder = _create_checkpoint_folder(saving, first) + _create_checkpoint_folder(saving, second) + + instruction = strategy.get_checkpoint_instruction(training_progress=first) + assert instruction.save_current + instruction = strategy.get_checkpoint_instruction(training_progress=second) + assert instruction.checkpoints_to_delete, "strategy should ask for the older checkpoint to go" + + # Deleting the evicted checkpoint must now succeed rather than raise OSError. + saving._delete_checkpoint(training_progress=first) + assert not first_folder.exists() diff --git a/tests/models/nemotron/test_mamba2_mixer.py b/tests/models/nemotron/test_mamba2_mixer.py index 0c8bff9b9..b9d44ac5b 100644 --- a/tests/models/nemotron/test_mamba2_mixer.py +++ b/tests/models/nemotron/test_mamba2_mixer.py @@ -152,8 +152,10 @@ def test_reset_parameters_is_a_noop_on_meta_device(): @pytest.mark.skipif(HAS_FUSED_KERNELS, reason="mamba-ssm and causal-conv1d are installed") def test_fused_backend_raises_a_helpful_error_when_kernels_are_missing(): # Requesting a backend whose dependencies are absent must fail loudly at construction time - # rather than silently falling back to a slower path. - with pytest.raises(ValueError, match="requires mamba-ssm and causal-conv1d"): + # rather than silently falling back to a slower path. The message names only the packages that + # are actually missing, so assert on the stable part rather than on a fixed package list - + # exactly one of the two extras may be installed. + with pytest.raises(ValueError, match=r"ssd_backend='fused' requires .* to be installed"): _make_mixer(ssd_backend=SSDBackend.FUSED)