diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 43d0c6e2d..412de3939 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,62 @@ 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 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 + 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_lorem_ipsum_nemotron_nano_fsdp2.yaml b/config_files/training/config_lorem_ipsum_nemotron_nano_fsdp2.yaml new file mode 100644 index 000000000..d571cdd74 --- /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. 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 + 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/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..80614905a --- /dev/null +++ b/config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml @@ -0,0 +1,423 @@ +# 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 `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 +# 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..9cef4ee40 --- /dev/null +++ b/docs/components/nemotron.md @@ -0,0 +1,243 @@ +# 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)). + +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 + +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) | + +## 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/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/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/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..62f78d7d8 --- /dev/null +++ b/src/modalities/models/components/mamba2/mamba2_mixer.py @@ -0,0 +1,357 @@ +# 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 +(``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..c201a71b8 --- /dev/null +++ b/src/modalities/models/components/mamba2/ssd.py @@ -0,0 +1,341 @@ +# 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 +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..03798263a --- /dev/null +++ b/src/modalities/models/components/moe/experts.py @@ -0,0 +1,182 @@ +# 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 +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..80c93fab4 --- /dev/null +++ b/src/modalities/models/components/moe/load_balancing.py @@ -0,0 +1,207 @@ +# 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" +(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..88ac07bb2 --- /dev/null +++ b/src/modalities/models/components/moe/moe.py @@ -0,0 +1,167 @@ +# 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 +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..9902b4d73 --- /dev/null +++ b/src/modalities/models/components/moe/router.py @@ -0,0 +1,138 @@ +# 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 +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..9d9dd0a41 --- /dev/null +++ b/src/modalities/models/nemotron/layer_pattern.py @@ -0,0 +1,95 @@ +# 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 +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..76cf3fdca --- /dev/null +++ b/src/modalities/models/nemotron/nemotron_attention.py @@ -0,0 +1,180 @@ +# 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 +``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..eb40ae893 --- /dev/null +++ b/src/modalities/models/nemotron/nemotron_layer_specs.py @@ -0,0 +1,352 @@ +# 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 +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..73f670dc4 --- /dev/null +++ b/src/modalities/models/nemotron/nemotron_layers.py @@ -0,0 +1,135 @@ +# 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 +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..721f09420 --- /dev/null +++ b/src/modalities/models/nemotron/nemotron_mlp.py @@ -0,0 +1,47 @@ +# 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 +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..d7947735c --- /dev/null +++ b/src/modalities/models/nemotron/nemotron_model.py @@ -0,0 +1,334 @@ +# 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, +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..5e9a75d5c --- /dev/null +++ b/src/modalities/models/nemotron/nemotron_stages_generator.py @@ -0,0 +1,250 @@ +# 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 +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..4b8461f9a --- /dev/null +++ b/src/modalities/utils/nemotron_mfu.py @@ -0,0 +1,194 @@ +"""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 | 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. + """ + + 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: 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 + + # 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: Optional[int] = None, + world_size: int = 1, + model_parts: torch.nn.Module | list[torch.nn.Module] = None, + 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 | 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. + """ + 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( + 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, device_mesh # 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. + + 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 | 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 + + 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 + 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/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/fsdp2_parallelization/nemotron_fsdp2_config.yaml b/tests/fsdp2_parallelization/nemotron_fsdp2_config.yaml new file mode 100644 index 000000000..811d6bd66 --- /dev/null +++ b/tests/fsdp2_parallelization/nemotron_fsdp2_config.yaml @@ -0,0 +1,165 @@ +# 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. + +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..b9d44ac5b --- /dev/null +++ b/tests/models/nemotron/test_mamba2_mixer.py @@ -0,0 +1,238 @@ +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. 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) + + +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..cf239ace3 --- /dev/null +++ b/tests/models/nemotron/test_nemotron_stages_and_mfu.py @@ -0,0 +1,392 @@ +"""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 + + +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 diff --git a/tests/models/nemotron/test_ssd.py b/tests/models/nemotron/test_ssd.py new file mode 100644 index 000000000..a7392db32 --- /dev/null +++ b/tests/models/nemotron/test_ssd.py @@ -0,0 +1,271 @@ +# `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 +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..617f1c690 --- /dev/null +++ b/tests/test_yaml_configs/nemotron_config_initialization.yaml @@ -0,0 +1,88 @@ +# 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 + 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