From 67adcd4854be45c571d7e4f2a0e4a0edce0f3f3e Mon Sep 17 00:00:00 2001 From: Max Luebbering <2804731+le1nux@users.noreply.github.com> Date: Mon, 7 Sep 2026 11:46:32 +0200 Subject: [PATCH] feat: added expert parallelism to nemotron-3 nano architecture among other improvements --- .gitignore | 3 +- ...g_fineweb_nemotron_nano_30b_a3b_fsdp2.yaml | 662 ++++++++++++++++ ..._nemotron_nano_30b_a3b_fsdp2_compiled.yaml | 734 ++++++++++++++++++ ...config_fineweb_nemotron_nano_ep_fsdp2.yaml | 609 +++++++++++++++ .../config_fineweb_nemotron_nano_fsdp2.yaml | 2 +- src/modalities/config/config.py | 45 +- src/modalities/loss_functions.py | 221 +++++- .../models/components/moe/experts.py | 45 +- src/modalities/models/gpt2/gpt2_model.py | 19 + src/modalities/models/model_factory.py | 88 ++- .../models/parallelism/expert_parallelism.py | 287 +++++++ src/modalities/registry/components.py | 15 +- .../running_env/fsdp/device_mesh.py | 110 ++- .../fsdp_gradient_clipper.py | 99 ++- tests/config/test_device_mesh_config.py | 50 ++ .../nemotron_ep_fsdp2_config.yaml | 182 +++++ .../test_expert_parallelism_fsdp2.py | 247 ++++++ .../nemotron/test_expert_parallelism.py | 106 +++ 18 files changed, 3451 insertions(+), 73 deletions(-) create mode 100644 config_files/training/config_fineweb_nemotron_nano_30b_a3b_fsdp2.yaml create mode 100644 config_files/training/config_fineweb_nemotron_nano_30b_a3b_fsdp2_compiled.yaml create mode 100644 config_files/training/config_fineweb_nemotron_nano_ep_fsdp2.yaml create mode 100644 src/modalities/models/parallelism/expert_parallelism.py create mode 100644 tests/config/test_device_mesh_config.py create mode 100644 tests/fsdp2_parallelization/nemotron_ep_fsdp2_config.yaml create mode 100644 tests/fsdp2_parallelization/test_expert_parallelism_fsdp2.py create mode 100644 tests/models/nemotron/test_expert_parallelism.py diff --git a/.gitignore b/.gitignore index d2c00d089..254ff9967 100644 --- a/.gitignore +++ b/.gitignore @@ -175,4 +175,5 @@ tutorials/scaling_up/experiments_old/* results/* tutorials/einsum_transformer/experiments/* tutorials/warmstart/experiments/* - +data/experiments/* +soofi/* \ No newline at end of file diff --git a/config_files/training/config_fineweb_nemotron_nano_30b_a3b_fsdp2.yaml b/config_files/training/config_fineweb_nemotron_nano_30b_a3b_fsdp2.yaml new file mode 100644 index 000000000..b55b3a685 --- /dev/null +++ b/config_files/training/config_fineweb_nemotron_nano_30b_a3b_fsdp2.yaml @@ -0,0 +1,662 @@ +# Nemotron-3 Nano 30B-A3B - FULL architecture, real-data (FineWeb) training config. +# +# This is the full-depth counterpart to config_fineweb_nemotron_nano_fsdp2.yaml. That file trims the +# depth from 52 to 16 layers so it fits four GPUs; this one keeps the published architecture intact +# and takes its data, optimizer and run plumbing from the 16-layer file: +# +# architecture identical to config_nemotron3_nano_30b_a3b_fsdp2.yaml (see DEVIATIONS below) +# data the Llama-3 tokenized FineWeb slices used by config_fineweb_nemotron_nano_fsdp2.yaml +# kernels ssd_backend: fused, experts_backend: grouped_mm +# +# ARCHITECTURE (from the model report arXiv:2512.20848 Table 1 / Figure 2, 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 +# +# Parameter counts, measured by instantiating this architecture on a meta device: +# +# total 31.563B (23 x 1297.5M MoE = 29.842B, 23 x 38.7M Mamba = 0.891B, +# 6 x 23.4M attention = 0.140B, 0.690B embedding + untied lm_head) +# active per token 3.565B (top-6 of 128 routed + both shared experts, all dense layers, +# and the 0.690B embedding + lm_head, per +# NemotronMFUCalculator.count_active_parameters) +# +# DEVIATIONS FROM config_nemotron3_nano_30b_a3b_fsdp2.yaml +# Two settings differ, both because that file is a lorem-ipsum smoke config and this one runs on real +# data at scale. The 52-layer pattern, d_model, every layer spec and every optimizer value are +# byte-identical. +# +# 1. vocab_size 131072 -> 128256. Dictated by the data, not by the architecture: the FineWeb +# slices are tokenized with the Llama-3 tokenizer (128256 tokens, already a multiple of 128). +# Keeping 131072 would still run - no token id exceeds 128255 - but 2816 embedding rows would +# never receive a gradient while their lm_head rows only ever get pushed down, i.e. 15.1M dead +# parameters. Restore 131072 only together with data tokenized by the Nemotron tokenizer. +# 2. ssd_backend native -> fused, matching the note in that file ("for real runs install the +# optional extra and switch to ssd_backend=fused"). Requires `pip install -e '.[mamba]'`. The +# two backends are numerically equivalent to 1e-3 (tests/models/components/test_mamba2.py). +# +# The step profile (gradient_accumulation_steps), schedule length (warmup_steps) and the interval / +# checkpoint-retention settings are not architecture; they are set for this data and hardware and are +# documented individually below. +# +# COMPILATION +# NOTHING IS COMPILED IN THIS CONFIG. Everything runs eager. Two optional compilations exist and are +# documented where they would go: +# +# the CE numeric core see the comment above `loss_fn`. Cheap, and worth turning on if the +# vocabulary projection puts you out of memory (7.83 -> 3.91 GiB measured). +# the model blocks see config_fineweb_nemotron_nano_30b_a3b_fsdp2_compiled.yaml, which is +# this file plus per-block torch.compile. +# +# The model blocks are left uncompiled here because the two block types that hold 97% of the +# parameters are the two that do not compile cleanly. Probed on this machine (A100, torch +# 2.9.1+cu128, bf16, one block of each type, fused/grouped_mm backends as configured here): +# +# block fullgraph=True graph breaks fullgraph=False +# NemotronAttentionLayer OK 0 OK, bit-exact +# NemotronMoELayer FAILS: aten.bincount 3 OK, bit-exact +# Mamba2Layer (fused) FAILS: causal_conv1d_cuda 11 OK, bit-exact +# +# ("bit-exact" is per block, in isolation. It does not compose to a bit-identical model: discrete +# top-k routing flips for a few borderline tokens once bf16 differences accumulate. See the +# COMPILATION section of config_fineweb_nemotron_nano_30b_a3b_fsdp2_compiled.yaml for measurements.) +# +# So adding a `model / compiled` component as llama_8B_fsdp2.yaml does would *crash*, since +# `CompiledModelConfig.fullgraph` defaults to True. Setting `fullgraph: false` runs, but is not worth +# it: the Mamba block shatters into 12 graphs of ~46 ops +# (Dynamo cannot trace the causal-conv1d pybind extension or the Triton autotuner's +# `best_config.kwargs[...]` lookup), and the MoE block into 4, with the break landing immediately +# before the grouped matmul that dominates its cost. That buys nearly nothing for 46 of 52 layers +# while paying full compile latency. The 6 attention layers compile cleanly but are 0.14B of 31.563B +# parameters. +# +# `torch._dynamo.config.capture_dynamic_output_shape_ops = True` does let the MoE block reach +# fullgraph with 0 breaks, but it is a global Dynamo flag with no YAML knob (it would need a code +# change), it does not help the Mamba blocks at all, and it perturbs the block output by 5.3e-3 +# relative - so it is not enabled here. +# +# Verified as a precondition, since it would fail silently: the in-place `router.tokens_per_expert` +# buffer update inside the MoE forward, which the auxiliary-loss-free expert bias depends on, and +# `last_aux_loss` both survive compilation in every configuration above. +# +# DATA +# The same two byte-exact, provably disjoint slices of a document-shuffled FineWeb sample-10BT +# tokenized with the Llama-3 tokenizer (vocab 128256, 4 bytes per token) as the 16-layer config: +# +# role file tokens source token range +# train ..._train-head-2B-tokens.pbin 2,000,000,000 [0, 2e9) +# eval ..._eval-tail-2M-tokens.pbin 2,000,000 [9,692,110,008, 9,694,110,008) +# +# The split is a head/tail cut rather than a random split, so disjointness is a property of the byte +# offsets and needs no bookkeeping. The source file is document-shuffled (seed 42), so the head is +# distributionally representative. +# +# THE TOKEN BUDGET IS THE MAIN LIMITATION OF THIS CONFIG. 2B tokens is ~0.06 tokens per parameter - +# three orders of magnitude below a compute-optimal budget for a 31.6B / 3.6B-active model. One pass +# over the head is a scale and plumbing validation on real data, not a convergence run. For a longer +# run: +# +# - `fineweb_sample-10BT.pbin` (38.9 GB, ~9.7B tokens) is the unshuffled full sample. The shuffled +# variant this head was carved from *contains* the eval tail, so either re-carve a head slice or +# use a held-out set from a different corpus. +# - `fineweb_sample-100BT.pbin` (391 GB, ~97B tokens) is the largest slice available locally. +# CAUTION, NOT VERIFIED HERE: FineWeb's smaller sample subsets are documented as subsets of the +# larger ones, so the eval tail above is most likely *inside* sample-100BT. Check for overlap or +# hold out a different set before quoting an eval loss against it. +# +# At sequence length 8192 with `reuse_last_target: true` the slices yield 244,140 train and 244 eval +# sequences. A larger `..._eval-tail-20M-tokens.pbin` (2,441 sequences) is also available; it is a +# superset of the 2M slice, so the two are alternatives, not independent sets. The 2M slice is wired +# up here because it is ~8x cheaper per evaluation. Both are disjoint from the training head. +# +# SIZE, PARALLELISM AND MEMORY +# +# NOT YET RUN. Unlike config_fineweb_nemotron_nano_fsdp2.yaml, whose memory and loss figures are +# measured, every number in this section is arithmetic from the parameter count above. Treat it as a +# feasibility estimate and measure before trusting it. +# +# The model does not fit on one node. Persistent training state is 16 bytes per parameter (fp32 +# shard + fp32 grad + two fp32 Adam moments) = 470 GiB, sharded `world_size` ways by FSDP2 full +# shard, plus per GPU roughly: +# +# unsharded bf16 all-gather of the largest FSDP unit (one 1297.5M-parameter MoE layer, 2.4 GiB) +# plus its prefetch ~4.8 GiB +# checkpointed layer inputs, 52 x 8192 x 2688 x 2 B at micro batch 1 ~2.2 GiB +# cross-entropy over 8192 x 128256 logits, forward + backward. THIS ONE IS +# MEASURED on an A100, eager as configured here: 1.96 GiB of logits plus a +# 5.87 GiB transient, i.e. 7.83 GiB peak. Compiling the CE core cuts the +# transient to 1.96 GiB (3.91 GiB peak). Largest activation term either way. ~7.8 GiB +# +# world size persistent / GPU estimated peak / GPU +# 4 117.6 GiB IMPOSSIBLE -> exceeds an 80 GiB card on optimizer state +# alone, before a single activation +# 8 58.8 GiB ~74 GiB -> do not attempt on 80 GB cards +# 16 29.4 GiB ~44 GiB -> minimum; fits 80 GB, tight on 64 GB +# 32 14.7 GiB ~30 GiB -> recommended +# +# 16 GPUs is the floor and 32 is the comfortable target. `data_parallel_shard_degree: -1` shards over +# the whole world size; do not switch to HSDP (replicate across nodes) - the memory table above +# assumes no replication. Tensor and expert parallelism are not supported for this architecture yet. +# +# GLOBAL BATCH depends on the world size, since micro batch and accumulation are fixed here: +# +# world size tokens / step steps for one pass over the 2B head +# 16 1,048,576 1,907 +# 32 2,097,152 953 +# +# `warmup_steps: 100` below is ~5% of the 16-GPU row. It is an absolute step count, so re-tune it if +# you change the world size, the accumulation, or the token budget. +# +# A DCP checkpoint is ~353 GiB (fp32 params plus both fp32 Adam moments for 31.563B parameters). +# With `k: 1` up to ~706 GiB is transiently on disk while a new one is written next to the retained +# one. Check free space before lowering `checkpointing_interval_in_steps`. +# +# EXPECTED SANITY SIGNALS +# - initial loss ~= ln(128256) = 11.76. A different value means the vocabulary and the +# tokenization of the data disagree; this shows up on the very first step. +# - held-out loss should track the train loss at the same step. A noticeably *lower* eval loss +# would indicate the eval slice leaking into training. +# +# Run with (example: 4 nodes x 8 GPUs = 32; set MASTER_ADDR to the rank-0 host): +# torchrun --nnodes 4 --nproc_per_node 8 --node_rank $NODE_RANK \ +# --master_addr $MASTER_ADDR --master_port 29555 \ +# $(which modalities) run \ +# --config_file_path config_files/training/config_fineweb_nemotron_nano_30b_a3b_fsdp2.yaml \ +# --experiments_root_path data/experiments +# +# Use one virtual environment for both `torchrun` and `modalities` - mixing them picks up a different +# torch build and the fused Mamba kernels will be missing. +# +# NO THROUGHPUT OR MFU NUMBER IS QUOTED HERE because none has been measured for this configuration. +# Measure on idle GPUs before using MFU to tune anything. + +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: /raid/s3/opengptx/max_lue/repositories/training_datasets/processed/fineweb_sample-10BT-shuffled_train-head-2B-tokens.pbin + test_dataset_path: /raid/s3/opengptx/max_lue/repositories/training_datasets/processed/fineweb_sample-10BT-shuffled_eval-tail-2M-tokens.pbin + experiments_root_path: ${modalities_env:experiments_root_path} + intervals: + # Throughput, loss, MFU, LR and peak memory are printed to the console every interval by the + # trainer itself (not by a subscriber), so this value sets time-to-first-output. The logging + # block does a dist.barrier() and an all-reduce, which is negligible next to a step of this + # size, so 1 is kept here: at ~1M tokens per step there are only ~1,900 steps in the whole run. + training_log_interval_in_steps: 1 + # ~353 GiB per checkpoint, see the header. 500 steps is ~4 checkpoints over a 16-GPU pass; each + # write moves a third of a terabyte, so this is an I/O decision as much as a safety one. + checkpointing_interval_in_steps: 500 + evaluation_interval_in_steps: 250 + 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: + # Micro batch 1 is forced by the activation and logit memory at sequence length 8192; the global + # batch comes entirely from accumulation and the world size (see the header table). 8 accumulation + # steps give 1.05M tokens per step at 16 GPUs and 2.10M at 32, both sane for a model this size. + gradient_accumulation_steps: 8 + local_train_micro_batch_size: 1 + # The published training length, unchanged from config_nemotron3_nano_30b_a3b_fsdp2.yaml. The + # data is continuously packed, so any value works; 8192 / chunk_size 128 = 64 Mamba scan chunks. + 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} + # Pretraining: the last target token of a sample is the first input token of the next, so no + # token is wasted at block boundaries. + reuse_last_target: true + +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} + reuse_last_target: true + +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: + # Kept at 1 even though evaluation is forward-only: at sequence length 8192 the logits, not + # the activations, dominate (2.0 GiB in bf16 per sequence before the fp32 upcast), so a + # larger eval batch is the most likely place for this config to run out of memory. + # 244 eval sequences with drop_last means 15 steps per pass at 16 ranks, 7 at 32; the + # remainder is dropped, so the exact eval subset depends on the world size. + batch_size: 1 + 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: + # Keep only the newest checkpoint - each is ~353 GiB. Rotation with the `dcp` execution is + # fixed (see tests/checkpointing/test_dcp_checkpoint_deletion.py). + 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. +# +# The cross-entropy runs EAGER here. Compiling just its numeric core is available and cheap - it cut +# measured peak allocation for the CE forward+backward from 7.83 GiB to 3.91 GiB at this config's +# 8192 x 128256 logits (the transient drops 5.87 -> 1.96 GiB on top of 1.96 GiB of resident logits), +# for a bit-identical loss value - so if you hit an out-of-memory in the vocabulary projection, this +# is the first thing to turn on: +# +# compiled_clm_loss: # add as a top-level component +# component_key: loss +# variant_key: compiled +# config: +# backend: inductor +# loss: +# component_key: loss +# variant_key: clm_cross_entropy_loss +# config: +# target_key: ${settings.referencing_keys.target_key} +# prediction_key: ${settings.referencing_keys.prediction_key} +# +# and replace the first entry of `losses` below with a BY_REFERENCE to it. It must wrap the *inner* +# CE loss, not the outer `weighted_sum`: WeightedSumLoss does not implement `compile` and raises +# NotImplementedError. +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.01725, 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 1297.5M-parameter MoE block (2.4 GiB in bf16) is + # unsharded at a time. Raising this multiplies that all-gather. + 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} + # Llama-3 tokenizer, matching the tokenization of the FineWeb slices above. Already a multiple + # of 128 (128256 = 1002 * 128), so no padding is needed. This is the one architecture value that + # differs from config_nemotron3_nano_30b_a3b_fsdp2.yaml (131072); see DEVIATIONS in the header. + vocab_size: 128256 + # Full published depth. Keep this in sync with the length of layer_pattern - the model + # validates it. + n_layer: 52 + n_embd: 2688 + 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 + # 8192 / 128 = 64 scan chunks. + chunk_size: 128 + ssd_backend: fused + 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 + +# Warmup then cosine decay to 10% of peak, the standard LLM pretraining schedule and what the +# reference recipe uses. `warmup_steps` is absolute and must stay well below `total_steps`, which is +# derived from the dataset and therefore depends on the world size: ~1,907 steps at 16 GPUs, ~953 at +# 32. 100 is ~5% of the former and ~10% of the latter. Re-tune it if you change the world size, the +# accumulation, or the token budget. +lr_scheduler: + component_key: scheduler + variant_key: linear_warmup_cosine_annealing_lr + config: + optimizer: + instance_key: optimizer + pass_type: BY_REFERENCE + warmup_steps: 100 + total_steps: ${settings.training_target.num_target_steps} + initial_lr: 4.5e-5 + max_lr: 4.5e-4 + final_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: 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 + +# OFFLINE, since compute nodes for a run this size usually have no outbound network. Sync afterwards +# with `wandb sync wandb_storage/`. +evaluation_subscriber: + component_key: results_subscriber + variant_key: wandb + config: + global_rank: ${settings.cuda_env.global_rank} + project: modalities_nemotron_nano_30b_a3b_fineweb + 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_fineweb_nemotron_nano_30b_a3b_fsdp2_compiled.yaml b/config_files/training/config_fineweb_nemotron_nano_30b_a3b_fsdp2_compiled.yaml new file mode 100644 index 000000000..9217755b9 --- /dev/null +++ b/config_files/training/config_fineweb_nemotron_nano_30b_a3b_fsdp2_compiled.yaml @@ -0,0 +1,734 @@ +# Nemotron-3 Nano 30B-A3B - FULL architecture, real-data (FineWeb), WITH torch.compile. +# +# Identical to config_fineweb_nemotron_nano_30b_a3b_fsdp2.yaml in every respect except that the +# transformer blocks are additionally run through torch.compile (see `compiled_model`). Diff the two +# files to see exactly that: one new component, one changed reference in `fsdp_model`, this header. +# +# READ THE COMPILATION SECTION BELOW BEFORE USING THIS VARIANT. Compiling this architecture is +# measurably close to a no-op: the two block types holding 97% of the parameters cannot be captured +# into a single graph, so they end up as 12 and 4 fragments respectively. It is correct - all three +# block types were verified bit-exact against eager - but the expected speedup is small and the +# compile latency is not. The uncompiled sibling config is the recommended default; this file exists +# for measuring whether the fragmented compile is worth anything on your hardware. +# +# This is the full-depth counterpart to config_fineweb_nemotron_nano_fsdp2.yaml. That file trims the +# depth from 52 to 16 layers so it fits four GPUs; this one keeps the published architecture intact +# and takes its data, optimizer and run plumbing from the 16-layer file: +# +# architecture identical to config_nemotron3_nano_30b_a3b_fsdp2.yaml (see DEVIATIONS below) +# data the Llama-3 tokenized FineWeb slices used by config_fineweb_nemotron_nano_fsdp2.yaml +# kernels ssd_backend: fused, experts_backend: grouped_mm +# compile per-block torch.compile, fullgraph=false, plus the compiled CE core +# +# ARCHITECTURE (from the model report arXiv:2512.20848 Table 1 / Figure 2, 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 +# +# Parameter counts, measured by instantiating this architecture on a meta device: +# +# total 31.563B (23 x 1297.5M MoE = 29.842B, 23 x 38.7M Mamba = 0.891B, +# 6 x 23.4M attention = 0.140B, 0.690B embedding + untied lm_head) +# active per token 3.565B (top-6 of 128 routed + both shared experts, all dense layers, +# and the 0.690B embedding + lm_head, per +# NemotronMFUCalculator.count_active_parameters) +# +# DEVIATIONS FROM config_nemotron3_nano_30b_a3b_fsdp2.yaml +# Two settings differ, both because that file is a lorem-ipsum smoke config and this one runs on real +# data at scale. The 52-layer pattern, d_model, every layer spec and every optimizer value are +# byte-identical. +# +# 1. vocab_size 131072 -> 128256. Dictated by the data, not by the architecture: the FineWeb +# slices are tokenized with the Llama-3 tokenizer (128256 tokens, already a multiple of 128). +# Keeping 131072 would still run - no token id exceeds 128255 - but 2816 embedding rows would +# never receive a gradient while their lm_head rows only ever get pushed down, i.e. 15.1M dead +# parameters. Restore 131072 only together with data tokenized by the Nemotron tokenizer. +# 2. ssd_backend native -> fused, matching the note in that file ("for real runs install the +# optional extra and switch to ssd_backend=fused"). Requires `pip install -e '.[mamba]'`. The +# two backends are numerically equivalent to 1e-3 (tests/models/components/test_mamba2.py). +# +# The step profile (gradient_accumulation_steps), schedule length (warmup_steps) and the interval / +# checkpoint-retention settings are not architecture; they are set for this data and hardware and are +# documented individually below. +# +# COMPILATION +# Two things are compiled here: every transformer block (`compiled_model`) and the cross-entropy +# numeric core (`compiled_clm_loss`). They are independent; only the first distinguishes this file +# from its uncompiled sibling. +# +# WHAT WAS MEASURED. Each block type was probed individually on this machine (A100-SXM4-80GB, torch +# 2.9.1+cu128, bf16, with the fused / grouped_mm backends this config selects): +# +# block fullgraph=True graph breaks fullgraph=False +# NemotronAttentionLayer OK 0 OK, bit-exact +# NemotronMoELayer FAILS: aten.bincount 3 OK, bit-exact +# Mamba2Layer (fused) FAILS: causal_conv1d_cuda 11 OK, bit-exact +# +# "bit-exact" is a per-block statement: max|compiled - eager| was exactly 0 on each block's output in +# isolation. +# +# IT DOES NOT COMPOSE TO A BIT-IDENTICAL MODEL, and the reason is worth understanding before you +# compare two runs. Measured on a 7-layer "MEM*EME" scale model through the exact AC -> compile chain +# below, compiled against eager: +# +# MoE layer 0: 0 / 1024 routing slots differ +# MoE layer 1: 0 / 1024 routing slots differ +# MoE layer 2: 5 / 1024 routing slots differ <- 0.16% +# -> 3 of 512 tokens with max|dlogit| > 0.1; max|dlogit| 1.03 against a logit absmax of 2.81 +# -> mean|dlogit| 2.9e-3, aux loss delta 5.2e-9, total gradient norm delta 1.4e-5 (rel ~1.4e-4) +# +# Nothing is broken here. Top-k routing is a *discrete* decision, so once bf16 rounding differences +# accumulate through a few layers, a handful of borderline tokens cross the top-6 boundary and get a +# different expert set. Those few tokens then differ a lot, which is why max|dlogit| looks alarming +# while every aggregate quantity - mean logit difference, aux loss, gradient norm - is unchanged to +# 4-5 digits. The same sensitivity exists between any two numerically-different-but-equally-valid +# execution paths of an MoE model, compile or not. +# +# PRACTICAL CONSEQUENCE: do not expect this config to reproduce the uncompiled one step for step, and +# do not treat a small loss divergence between the two as evidence of a bug. Compare loss *curves*, +# not individual steps. With 52 layers instead of 7 the flip count will be higher than measured above. +# +# WHY fullgraph MUST BE false. `CompiledModelConfig.fullgraph` defaults to true, which raises on this +# architecture. The Mamba-2 block calls into the causal-conv1d pybind extension, which Dynamo refuses +# to trace ("function marked as skipped"), and additionally reads the Triton autotuner's +# `_state_passing_bwd_kernel.best_config.kwargs["BLOCK_SIZE"]`, an untraceable attribute lookup. The +# MoE block calls `torch.bincount`, whose output shape is data-dependent. Both are inherent to the +# fast kernels, not incidental. +# +# WHAT YOU ACTUALLY GET. With fullgraph false the breaks become fragmentation rather than errors: the +# Mamba block becomes 12 graphs over ~46 ops and the MoE block 4 graphs, with its break landing +# immediately before the grouped matmul that dominates its cost. The 6 attention layers compile into +# a single clean graph but are only 0.14B of 31.563B parameters. So 46 of 52 layers - 97% of the +# parameters - are compiled in fragments too small to fuse across the expensive operations. Expect a +# small gain at best, against a real one-time compile latency (three distinct block classes, so +# roughly three compilations, not 52). MEASURE IT: if `samples/s` at steady state does not beat the +# uncompiled config on idle GPUs, use the uncompiled one. +# +# NOT ENABLED: `torch._dynamo.config.capture_dynamic_output_shape_ops = True` does get the MoE block +# to fullgraph with 0 breaks, but it is a global Dynamo flag with no YAML knob (it needs a code +# change), it does nothing for the Mamba blocks, and it perturbs the MoE block output by 5.3e-3 +# relative - no longer bit-exact. +# +# VERIFIED AS A PRECONDITION, because it would otherwise fail silently: the in-place +# `router.tokens_per_expert` buffer update inside the MoE forward, which the auxiliary-loss-free +# expert bias in the optimizer depends on, and `last_aux_loss` both survive compilation in every +# configuration above. Compiled blocks also keep working with the initializer, which strips the +# `_orig_mod.` and `_checkpoint_wrapped_module.` FQN prefixes before matching its per-layer regexes +# (see initialization_routines.py and the regression guard in +# tests/models/nemotron/test_nemotron_initialization.py). +# +# The compiled CE core is the one unambiguous win and is present in both configs: at 8192 x 128256 +# logits, measured peak allocation for the CE forward + backward drops from 7.83 GiB to 3.91 GiB for +# a bit-identical loss. +# +# DATA +# The same two byte-exact, provably disjoint slices of a document-shuffled FineWeb sample-10BT +# tokenized with the Llama-3 tokenizer (vocab 128256, 4 bytes per token) as the 16-layer config: +# +# role file tokens source token range +# train ..._train-head-2B-tokens.pbin 2,000,000,000 [0, 2e9) +# eval ..._eval-tail-2M-tokens.pbin 2,000,000 [9,692,110,008, 9,694,110,008) +# +# The split is a head/tail cut rather than a random split, so disjointness is a property of the byte +# offsets and needs no bookkeeping. The source file is document-shuffled (seed 42), so the head is +# distributionally representative. +# +# THE TOKEN BUDGET IS THE MAIN LIMITATION OF THIS CONFIG. 2B tokens is ~0.06 tokens per parameter - +# three orders of magnitude below a compute-optimal budget for a 31.6B / 3.6B-active model. One pass +# over the head is a scale and plumbing validation on real data, not a convergence run. For a longer +# run: +# +# - `fineweb_sample-10BT.pbin` (38.9 GB, ~9.7B tokens) is the unshuffled full sample. The shuffled +# variant this head was carved from *contains* the eval tail, so either re-carve a head slice or +# use a held-out set from a different corpus. +# - `fineweb_sample-100BT.pbin` (391 GB, ~97B tokens) is the largest slice available locally. +# CAUTION, NOT VERIFIED HERE: FineWeb's smaller sample subsets are documented as subsets of the +# larger ones, so the eval tail above is most likely *inside* sample-100BT. Check for overlap or +# hold out a different set before quoting an eval loss against it. +# +# At sequence length 8192 with `reuse_last_target: true` the slices yield 244,140 train and 244 eval +# sequences. A larger `..._eval-tail-20M-tokens.pbin` (2,441 sequences) is also available; it is a +# superset of the 2M slice, so the two are alternatives, not independent sets. The 2M slice is wired +# up here because it is ~8x cheaper per evaluation. Both are disjoint from the training head. +# +# SIZE, PARALLELISM AND MEMORY +# +# NOT YET RUN. Unlike config_fineweb_nemotron_nano_fsdp2.yaml, whose memory and loss figures are +# measured, every number in this section is arithmetic from the parameter count above. Treat it as a +# feasibility estimate and measure before trusting it. +# +# The model does not fit on one node. Persistent training state is 16 bytes per parameter (fp32 +# shard + fp32 grad + two fp32 Adam moments) = 470 GiB, sharded `world_size` ways by FSDP2 full +# shard, plus per GPU roughly: +# +# unsharded bf16 all-gather of the largest FSDP unit (one 1297.5M-parameter MoE layer, 2.4 GiB) +# plus its prefetch ~4.8 GiB +# checkpointed layer inputs, 52 x 8192 x 2688 x 2 B at micro batch 1 ~2.2 GiB +# cross-entropy over 8192 x 128256 logits, forward + backward. THIS ONE IS +# MEASURED on an A100: 1.96 GiB of logits plus a 1.96 GiB transient with the +# compiled CE core configured below, i.e. 3.91 GiB peak; 7.83 GiB eager. +# Largest activation term either way. ~3.9 GiB +# +# world size persistent / GPU estimated peak / GPU +# 4 117.6 GiB IMPOSSIBLE -> exceeds an 80 GiB card on optimizer state +# alone, before a single activation +# 8 58.8 GiB ~70 GiB -> do not attempt on 80 GB cards +# 16 29.4 GiB ~40 GiB -> minimum; fits 80 GB, tight on 64 GB +# 32 14.7 GiB ~26 GiB -> recommended +# +# 16 GPUs is the floor and 32 is the comfortable target. `data_parallel_shard_degree: -1` shards over +# the whole world size; do not switch to HSDP (replicate across nodes) - the memory table above +# assumes no replication. Tensor and expert parallelism are not supported for this architecture yet. +# +# GLOBAL BATCH depends on the world size, since micro batch and accumulation are fixed here: +# +# world size tokens / step steps for one pass over the 2B head +# 16 1,048,576 1,907 +# 32 2,097,152 953 +# +# `warmup_steps: 100` below is ~5% of the 16-GPU row. It is an absolute step count, so re-tune it if +# you change the world size, the accumulation, or the token budget. +# +# A DCP checkpoint is ~353 GiB (fp32 params plus both fp32 Adam moments for 31.563B parameters). +# With `k: 1` up to ~706 GiB is transiently on disk while a new one is written next to the retained +# one. Check free space before lowering `checkpointing_interval_in_steps`. +# +# EXPECTED SANITY SIGNALS +# - initial loss ~= ln(128256) = 11.76. A different value means the vocabulary and the +# tokenization of the data disagree; this shows up on the very first step. +# - held-out loss should track the train loss at the same step. A noticeably *lower* eval loss +# would indicate the eval slice leaking into training. +# +# Run with (example: 4 nodes x 8 GPUs = 32; set MASTER_ADDR to the rank-0 host): +# torchrun --nnodes 4 --nproc_per_node 8 --node_rank $NODE_RANK \ +# --master_addr $MASTER_ADDR --master_port 29555 \ +# $(which modalities) run \ +# --config_file_path config_files/training/config_fineweb_nemotron_nano_30b_a3b_fsdp2_compiled.yaml \ +# --experiments_root_path data/experiments +# +# The W&B project is deliberately the same as the uncompiled config's, so the two runs land side by +# side and `samples/s` is directly comparable. Expect the first few steps to be dominated by compile +# latency; compare steady state, not step 1. +# +# Use one virtual environment for both `torchrun` and `modalities` - mixing them picks up a different +# torch build and the fused Mamba kernels will be missing. +# +# NO THROUGHPUT OR MFU NUMBER IS QUOTED HERE because none has been measured for this configuration. +# Measure on idle GPUs before using MFU to tune anything. + +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: /raid/s3/opengptx/max_lue/repositories/training_datasets/processed/fineweb_sample-10BT-shuffled_train-head-2B-tokens.pbin + test_dataset_path: /raid/s3/opengptx/max_lue/repositories/training_datasets/processed/fineweb_sample-10BT-shuffled_eval-tail-2M-tokens.pbin + experiments_root_path: ${modalities_env:experiments_root_path} + intervals: + # Throughput, loss, MFU, LR and peak memory are printed to the console every interval by the + # trainer itself (not by a subscriber), so this value sets time-to-first-output. The logging + # block does a dist.barrier() and an all-reduce, which is negligible next to a step of this + # size, so 1 is kept here: at ~1M tokens per step there are only ~1,900 steps in the whole run. + training_log_interval_in_steps: 1 + # ~353 GiB per checkpoint, see the header. 500 steps is ~4 checkpoints over a 16-GPU pass; each + # write moves a third of a terabyte, so this is an I/O decision as much as a safety one. + checkpointing_interval_in_steps: 500 + evaluation_interval_in_steps: 250 + 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: + # Micro batch 1 is forced by the activation and logit memory at sequence length 8192; the global + # batch comes entirely from accumulation and the world size (see the header table). 8 accumulation + # steps give 1.05M tokens per step at 16 GPUs and 2.10M at 32, both sane for a model this size. + gradient_accumulation_steps: 8 + local_train_micro_batch_size: 1 + # The published training length, unchanged from config_nemotron3_nano_30b_a3b_fsdp2.yaml. The + # data is continuously packed, so any value works; 8192 / chunk_size 128 = 64 Mamba scan chunks. + 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} + # Pretraining: the last target token of a sample is the first input token of the next, so no + # token is wasted at block boundaries. + reuse_last_target: true + +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} + reuse_last_target: true + +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: + # Kept at 1 even though evaluation is forward-only: at sequence length 8192 the logits, not + # the activations, dominate (2.0 GiB in bf16 per sequence before the fp32 upcast), so a + # larger eval batch is the most likely place for this config to run out of memory. + # 244 eval sequences with drop_last means 15 steps per pass at 16 ranks, 7 at 32; the + # remainder is dropped, so the exact eval subset depends on the world size. + batch_size: 1 + 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: + # Keep only the newest checkpoint - each is ~353 GiB. Rotation with the `dcp` execution is + # fixed (see tests/checkpointing/test_dcp_checkpoint_deletion.py). + 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} + +# The cross-entropy core is torch.compile'd. This is the one place in this architecture where +# compilation pays off, and it is the largest single activation term at this sequence length: with +# 8192 x 128256 logits, measured peak allocation for the CE forward+backward drops from 7.83 GiB to +# 3.91 GiB (A100, torch 2.9.1, bf16 logits) for a bit-identical loss value. Only the numeric +# tensor-in/tensor-out core is compiled, never the batch unpacking. +# +# NOTE: this must wrap the *inner* CE loss, not the outer `weighted_sum` - WeightedSumLoss does not +# implement `compile` and raises NotImplementedError. See COMPILATION in the header for why the +# model blocks themselves are deliberately left uncompiled. +compiled_clm_loss: + component_key: loss + variant_key: compiled + config: + backend: inductor + loss: + component_key: loss + variant_key: clm_cross_entropy_loss + config: + target_key: ${settings.referencing_keys.target_key} + prediction_key: ${settings.referencing_keys.prediction_key} + +# 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: + - instance_key: compiled_clm_loss + pass_type: BY_REFERENCE + - 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.01725, 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: compiled_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 1297.5M-parameter MoE block (2.4 GiB in bf16) is + # unsharded at a time. Raising this multiplies that all-gather. + layers_per_fsdp_unit: 1 + block_names: [Mamba2Layer, NemotronMoELayer, NemotronAttentionLayer, NemotronMLPLayer] + +# torch.compile, applied per transformer block AFTER activation checkpointing and BEFORE FSDP2 - +# the TorchTitan ordering. Activation checkpointing replaces each entry of `transformer.h` with a +# CheckpointWrapper, so the block that gets compiled here is the one *inside* the wrapper and the +# recompute in the backward pass runs the compiled graph. +# +# TWO SETTINGS HERE ARE LOAD-BEARING; see COMPILATION in the header for the measurements. +# +# fullgraph: false REQUIRED. The component default is true, which CRASHES on this architecture: +# the Mamba-2 blocks hit the causal-conv1d pybind extension and the MoE blocks +# hit `aten.bincount`, neither of which Dynamo can trace. With false, all three +# block types compile, each bit-exact against eager in isolation. The assembled +# model is NOT bit-identical, because discrete top-k routing flips for a few +# borderline tokens - see COMPILATION in the header. +# +# block_names MUST NOT list NemotronMLPLayer, even though `fsdp_model` above does. This +# layer_pattern contains no "-" layers, so no NemotronMLPLayer instance exists, +# and `get_compiled_model` raises ValueError for a block name that matches no +# module (fsdp2_wrapped tolerates the same name silently). Add it back only if +# you put "-" layers in the pattern. +compiled_model: + component_key: model + variant_key: compiled + config: + model: + instance_key: activation_checkpointed_model + pass_type: BY_REFERENCE + block_names: [Mamba2Layer, NemotronMoELayer, NemotronAttentionLayer] + fullgraph: false + debug: false + +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} + # Llama-3 tokenizer, matching the tokenization of the FineWeb slices above. Already a multiple + # of 128 (128256 = 1002 * 128), so no padding is needed. This is the one architecture value that + # differs from config_nemotron3_nano_30b_a3b_fsdp2.yaml (131072); see DEVIATIONS in the header. + vocab_size: 128256 + # Full published depth. Keep this in sync with the length of layer_pattern - the model + # validates it. + n_layer: 52 + n_embd: 2688 + 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 + # 8192 / 128 = 64 scan chunks. + chunk_size: 128 + ssd_backend: fused + 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 + +# Warmup then cosine decay to 10% of peak, the standard LLM pretraining schedule and what the +# reference recipe uses. `warmup_steps` is absolute and must stay well below `total_steps`, which is +# derived from the dataset and therefore depends on the world size: ~1,907 steps at 16 GPUs, ~953 at +# 32. 100 is ~5% of the former and ~10% of the latter. Re-tune it if you change the world size, the +# accumulation, or the token budget. +lr_scheduler: + component_key: scheduler + variant_key: linear_warmup_cosine_annealing_lr + config: + optimizer: + instance_key: optimizer + pass_type: BY_REFERENCE + warmup_steps: 100 + total_steps: ${settings.training_target.num_target_steps} + initial_lr: 4.5e-5 + max_lr: 4.5e-4 + final_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: 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 + +# OFFLINE, since compute nodes for a run this size usually have no outbound network. Sync afterwards +# with `wandb sync wandb_storage/`. +evaluation_subscriber: + component_key: results_subscriber + variant_key: wandb + config: + global_rank: ${settings.cuda_env.global_rank} + project: modalities_nemotron_nano_30b_a3b_fineweb + 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_fineweb_nemotron_nano_ep_fsdp2.yaml b/config_files/training/config_fineweb_nemotron_nano_ep_fsdp2.yaml new file mode 100644 index 000000000..8543dc84c --- /dev/null +++ b/config_files/training/config_fineweb_nemotron_nano_ep_fsdp2.yaml @@ -0,0 +1,609 @@ +# Nemotron-3 Nano real-data training config for 4x A100 80GB, WITH EXPERT PARALLELISM. +# +# Byte-identical to config_fineweb_nemotron_nano_fsdp2.yaml except for three things, so that the two +# can be compared directly as an expert-parallelism A/B: +# +# 1. `expert_parallel_degree: 4` in the device_mesh component. +# 2. A new `expert_parallelized_model` component between `model_raw` and +# `activation_checkpointed_model`. +# 3. This header. +# +# EXPERT PARALLELISM +# Without it, FSDP2 shards the 128 routed experts of each MoE layer over the 4 ranks and all-gathers +# the full stack again on every forward and backward pass. Each of the 7 MoE layers of this config +# holds 128 x (1856 x 2688 + 2688 x 1856) = 1.277B expert parameters, i.e. 2.38 GiB in bf16, so a +# step moves ~16.7 GiB of expert weights per rank in the forward pass alone - and activation +# checkpointing makes the forward run twice. +# +# With expert parallelism each rank instead *owns* 32 of the 128 experts outright and never gathers +# the other 96. What travels instead are the tokens: each rank routes its own tokens, sends them to +# the rank owning their expert (dispatch all-to-all), and gets the results back (combine all-to-all). +# At micro batch 1 x 4096 tokens and top-6 routing that is 4096 x 6 x 2688 x 2 B = 126 MiB per +# direction per layer, of which 3/4 leaves the rank - roughly an order of magnitude less traffic than +# the weight all-gather it replaces, and on NVLink rather than competing with it. +# +# The trade is that the expert matmuls get *smaller* per rank in one dimension and *larger* in +# another: each rank now evaluates 32 experts over tokens gathered from all 4 ranks, instead of 128 +# experts over its own tokens. Total expert FLOPs per rank are unchanged; the grouped matmul just +# sees 4x more rows per expert, which on A100 is the more efficient shape. +# +# MEASURED, 4x A100-SXM4-80GB (NVSwitch, NV12 between every pair), torch 2.9.1+cu128, this config +# versus config_fineweb_nemotron_nano_fsdp2.yaml, median over steps 4+ of a ~5 minute run: +# +# expert_parallel_degree samples/s tokens/s/GPU MFU peak mem speedup +# 1 (the non-EP config) 5.40 5530 0.170 48.97 GiB 1.000x +# 2 5.90 6042 0.190 52.42 GiB 1.093x +# 4 (this config) 7.10 7270 0.230 48.02 GiB 1.315x +# +# Per-step samples/s ranges were [5.30, 5.60], [5.80, 6.00] and [7.00, 7.30], i.e. non-overlapping. +# +# Two things are worth reading off this table. First, degree 2 is the worst of both worlds on memory: +# the experts are still FSDP-sharded (over the 2 remaining data-parallel ranks) so the all-gather +# buffer is still allocated, *and* the all-to-all buffers are added on top. Second, the speedup is +# not the MoE math getting faster. Benchmarked in isolation without FSDP, one expert-parallel MoE +# layer at these dimensions is 0.89x on the forward pass and 1.06x on forward+backward versus a +# replicated one - roughly break-even, because the 1.61 ms all-to-all pair (17.5% of the layer +# forward) costs about what the better grouped-matmul shape saves. The entire end-to-end win comes +# from no longer all-gathering expert weights. +# +# See src/modalities/models/parallelism/expert_parallelism.py for the dispatch/combine implementation. +# +# ARCHITECTURE +# 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 plus 2 shared experts, 32 query / 2 KV +# attention heads of head dim 128 - with depth cut from 52 to 16 layers to fit four GPUs. The layer +# pattern is the first 16 characters of the published pattern, preserving 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. +# +# DATA +# Both files are byte-exact, provably disjoint slices of a document-shuffled FineWeb sample-10BT +# tokenized with the Llama-3 tokenizer (vocab 128256, 4 bytes per token): +# +# role file tokens source token range +# train ..._train-head-2B-tokens.pbin 2,000,000,000 [0, 2e9) +# eval ..._eval-tail-2M-tokens.pbin 2,000,000 [9,692,110,008, 9,694,110,008) +# +# The split is a head/tail cut rather than a random split so that disjointness is a property of the +# byte offsets and needs no bookkeeping. The source file is document-shuffled (seed 42), so the head +# is distributionally representative. +# +# A larger 20M-token slice (`..._eval-tail-20M-tokens.pbin`, 4,882 sequences) covering tokens +# [9,674,110,008, 9,694,110,008) is also available for a final evaluation. Both eval slices are tails +# of the same file, so the 2M slice is a *subset* of the 20M one - they are alternatives, not +# independent sets. The 2M slice is wired up here because a full pass over the 20M one is 610 eval +# steps; the 2M slice is 61. Both are disjoint from the training head. +# +# The 2B-token head is sized to what this hardware can actually consume, not to the full 9.69B-token +# sample. To train on the full sample instead, point `train_dataset_path` at the original +# `fineweb_sample-10BT-shuffled-seed_42-bs_100.pbin` - but note that it *contains* the eval tail, so +# either re-carve a head slice or switch to a held-out set from a different corpus. +# +# SIZE AND MEMORY (4x A100-SXM4-80GB, FSDP2 full shard, AdamW, bf16 params / fp32 reduce) +# +# parameters 10.09B (16 layers; the larger vocab adds 419M over the GPT-2 vocab variant) +# active per token 1.57B (includes the 690M embedding + untied lm_head) +# global batch 32 sequences = 131,072 tokens (4 ranks x 1 micro batch x 8 accum steps) +# target steps 15,258 for one pass over the 2B-token head +# +# peak memory 49.0 GiB / GPU (measured, see VERIFICATION) +# +# Persistent optimizer state dominates that peak: fp32 shard + fp32 grad + two fp32 Adam moments is +# 16 bytes per parameter, sharded four ways, i.e. 37.6 GiB per GPU. The remaining ~11 GiB is the +# unsharded bf16 all-gather of the largest FSDP unit (one ~1.3B-parameter MoE layer, ~2.6 GiB) plus +# its prefetch, the checkpointed layer inputs (~350 MiB), and the vocabulary projection - 4096 x +# 128256 logits are ~1 GiB in bf16 and the cross-entropy upcasts them. +# +# For 40GB cards set `n_layer: 9` with `layer_pattern: "MEMEM*EME"`: 6.06B parameters, 22.6 GiB of +# persistent state, and a measured peak of 33.8 GiB. +# +# VERIFICATION +# Both variants were run end to end on the real data above (4x A100-SXM4-80GB): +# +# variant steps initial loss final loss held-out loss peak memory exit +# 16 layer 6 11.95 8.99 8.99 at step 5 49.0 GiB 0 +# 9 layer 40 11.94 6.96 7.38 at step 20 33.8 GiB 0 +# +# The initial loss matching ln(128256) = 11.76 is the check that the vocabulary and the tokenization +# of the data agree - a mismatched vocab shows up here immediately. The held-out loss tracking the +# train loss at the same step (8.99 vs 8.99; 7.38 vs 7.35) is the check on the data split: a +# noticeably *lower* eval loss would indicate the eval slice leaking into training. +# +# Those two runs used this file's architecture, data, optimizer and parallelism unchanged. They +# differed only in run length: the token budget was capped to the step count shown, `warmup_steps` +# was scaled down to match, and checkpointing was disabled so that no 113 GiB write was triggered. +# The 9-layer row additionally used the reduced depth described above. Neither run is a substitute +# for a convergence run - they establish that the configuration is valid, fits, and learns. +# +# 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_fineweb_nemotron_nano_fsdp2.yaml \ +# --experiments_root_path data/experiments +# +# Use one virtual environment for both `torchrun` and `modalities` - mixing them picks up a different +# torch build and the fused Mamba kernels will be missing. +# +# `ssd_backend: fused` requires the optional extra: `pip install -e '.[mamba]'`. The fused and native +# backends are numerically equivalent (verified to 1e-3 in tests/models/components/test_mamba2.py); +# switch to `native` if the kernels are unavailable, at a substantial throughput cost. +# +# NO THROUGHPUT NUMBER IS QUOTED HERE ON PURPOSE. The verification runs for this config shared their +# GPUs with an unrelated job at 100% utilization, so the samples/s and MFU they reported are measures +# of contention, not of this configuration. Measure on idle GPUs before using MFU to tune anything. + +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: /raid/s3/opengptx/max_lue/repositories/training_datasets/processed/fineweb_sample-10BT-shuffled_train-head-2B-tokens.pbin + test_dataset_path: /raid/s3/opengptx/max_lue/repositories/training_datasets/processed/fineweb_sample-10BT-shuffled_eval-tail-2M-tokens.pbin + experiments_root_path: ${modalities_env:experiments_root_path} + intervals: + # Throughput, loss, MFU, LR and peak memory are printed to the console every interval by the + # trainer itself (not by a subscriber), so this value sets time-to-first-output. At ~16 s/step + # a value of 10 means 2.7 minutes of apparent silence at startup, which reads like a hang; 5 + # halves that. The logging block does a dist.barrier() and an all-reduce, so do not set this to + # 1 for a long run - use 1 only for a smoke check where you want a line every step. + training_log_interval_in_steps: 1 + # A DCP checkpoint of this model is ~113 GiB (10.09B parameters plus fp32 Adam state, written by + # all 4 ranks). Combined with `k: 1` below, at most ~226 GiB is on disk transiently while a new + # checkpoint is written next to the retained one. Check free space before lowering this. + checkpointing_interval_in_steps: 5000 + evaluation_interval_in_steps: 250 + 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 buys a sane global batch (32 sequences / 131k tokens) at micro batch 1, + # which is what the ~40 GiB of resident optimizer state leaves room for at sequence length 4096. + gradient_accumulation_steps: 8 + local_train_micro_batch_size: 1 + sequence_length: 4096 + 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} + # Pretraining: the last target token of a sample is the first input token of the next, so no + # token is wasted at block boundaries. + reuse_last_target: true + +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} + reuse_last_target: true + +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: + # Forward-only, so a larger batch than training fits. 488 eval sequences across 4 ranks at + # batch size 2 is exactly 61 eval steps per pass. + batch_size: 2 + 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: + # Keep only the newest checkpoint - each is ~113 GiB. Rotation with the `dcp` execution is + # fixed (see tests/checkpointing/test_dcp_checkpoint_deletion.py). + 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 + # Expert parallelism is carved out of data_parallel_shard_degree, not multiplied into the world + # size: the 4 ranks still each process their own data shard, they just own 32 of the 128 routed + # experts each and exchange tokens by all-to-all. `dp_degree` below therefore still resolves to + # 4, and the global batch is unchanged from the non-EP config. + # + # At expert_parallel_degree == data_parallel_shard_degree (4 here) nothing is left of the + # data-parallel dimension for FSDP to shard the experts over, so expert weights are never + # all-gathered. Set this to 2 to shard experts 2-way and FSDP-shard each half over the + # remaining 2 ranks. + expert_parallel_degree: 4 + 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] + +expert_parallelized_model: + component_key: model + variant_key: expert_parallelized + config: + model: + instance_key: model_raw + pass_type: BY_REFERENCE + device_mesh: + instance_key: device_mesh + pass_type: BY_REFERENCE + +activation_checkpointed_model: + component_key: model + variant_key: activation_checkpointed + config: + model: + instance_key: expert_parallelized_model + 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} + # Llama-3 tokenizer, matching the tokenization of the FineWeb slices above. Already a multiple + # of 128 (128256 = 1002 * 128), so no padding is needed. + vocab_size: 128256 + 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 + # 4096 / 128 = 32 scan chunks. + chunk_size: 128 + ssd_backend: fused + 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 + +# Warmup then cosine decay to 10% of peak, the standard LLM pretraining schedule and what the +# reference recipe uses. `warmup_steps` is absolute, so it must stay well below `total_steps` +# (~15,258 here); 750 is ~5% of the run. Re-tune it if you change the token budget or batch size. +lr_scheduler: + component_key: scheduler + variant_key: linear_warmup_cosine_annealing_lr + config: + optimizer: + instance_key: optimizer + pass_type: BY_REFERENCE + warmup_steps: 750 + total_steps: ${settings.training_target.num_target_steps} + initial_lr: 4.5e-5 + max_lr: 4.5e-4 + final_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: 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_fineweb + 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_fineweb_nemotron_nano_fsdp2.yaml b/config_files/training/config_fineweb_nemotron_nano_fsdp2.yaml index 4aa913ba7..93ac36f8f 100644 --- a/config_files/training/config_fineweb_nemotron_nano_fsdp2.yaml +++ b/config_files/training/config_fineweb_nemotron_nano_fsdp2.yaml @@ -113,7 +113,7 @@ settings: # a value of 10 means 2.7 minutes of apparent silence at startup, which reads like a hang; 5 # halves that. The logging block does a dist.barrier() and an all-reduce, so do not set this to # 1 for a long run - use 1 only for a smoke check where you want a line every step. - training_log_interval_in_steps: 5 + training_log_interval_in_steps: 1 # A DCP checkpoint of this model is ~113 GiB (10.09B parameters plus fp32 Adam state, written by # all 4 ranks). Combined with `k: 1` below, at most ~226 GiB is on disk transiently while a new # checkpoint is written next to the retained one. Check free space before lowering this. diff --git a/src/modalities/config/config.py b/src/modalities/config/config.py index 0ae12f0ba..9cc33c470 100644 --- a/src/modalities/config/config.py +++ b/src/modalities/config/config.py @@ -24,6 +24,7 @@ PydanticFSDP1ModuleType, PydanticFSDP2ModuleType, PydanticLLMDataLoaderIFType, + PydanticLossIFType, PydanticLRSchedulerIFType, PydanticModelInitializationIFType, PydanticOptimizerIFType, @@ -41,7 +42,7 @@ PyTorchDtypes, has_bfloat_support, ) -from modalities.running_env.fsdp.device_mesh import ParallelismDegrees +from modalities.running_env.fsdp.device_mesh import ParallelismDegrees, has_parallelism_method from modalities.training.activation_checkpointing.activation_checkpointing_variants import ( ActivationCheckpointingVariants, ) @@ -85,6 +86,25 @@ class CLMCrossEntropyLossConfig(BaseModel): prediction_key: str +class ChunkedCLMCrossEntropyLossConfig(BaseModel): + # BY_REFERENCE to the (already wrapped) model; the loss borrows its lm_head and + # switches the model into skip_lm_head mode. See ChunkedCLMCrossEntropyLoss. + model: PydanticPytorchModuleType + target_key: str + prediction_key: str + num_chunks: Annotated[int, Field(strict=True, ge=1)] = 8 + + # avoid pydantic warning about the protected 'model_' namespace + model_config = ConfigDict(protected_namespaces=()) + + +class CompiledLossConfig(BaseModel): + # Wraps a raw loss (BY_REFERENCE) and compiles its tensor core in place, + # mirroring the CompiledModelConfig / model "compiled" variant. + loss: PydanticLossIFType + backend: str = "inductor" + + # Checkpointing class SaveEveryKStepsCheckpointingStrategyConfig(BaseModel): k: PositiveInt @@ -295,6 +315,9 @@ class FSDP2WrappedModelConfig(BaseModel): reshard_after_forward: bool = True device_mesh: PydanticDeviceMeshIFType layers_per_fsdp_unit: int = 1 + # Shard the lm_head as its own FSDP unit. Required when the model is trained with + # ChunkedCLMCrossEntropyLoss (the head is applied outside the model forward). + separate_lm_head_fsdp_unit: bool = False @model_validator(mode="after") def validate_mixed_precision_settings(self): @@ -309,11 +332,29 @@ def validate_mixed_precision_settings(self): def validate_dp_mesh_existence(self): if self.device_mesh.mesh_dim_names is None: raise ValueError(f"Device mesh {self.device_mesh=} has no defined mesh_dim_names.") - if ParallelismDegrees.DP_SHARD.value not in self.device_mesh.mesh_dim_names: + # Resolved via has_parallelism_method rather than mesh_dim_names because under expert + # parallelism dp_shard is a flattened dimension, which is addressable but not named. + if not has_parallelism_method(self.device_mesh, ParallelismDegrees.DP_SHARD): raise ValueError(f"Data parallelism key '{ParallelismDegrees.DP_SHARD.value}' not in {self.device_mesh=}") return self +class ExpertParallelizedModelConfig(BaseModel): + model: PydanticPytorchModuleOrListType + device_mesh: PydanticDeviceMeshIFType + + @model_validator(mode="after") + def validate_ep_mesh_existence(self) -> "ExpertParallelizedModelConfig": + if self.device_mesh.mesh_dim_names is None: + raise ValueError(f"Device mesh {self.device_mesh=} has no defined mesh_dim_names.") + if not has_parallelism_method(self.device_mesh, ParallelismDegrees.EP): + raise ValueError( + f"Expert parallelism key '{ParallelismDegrees.EP.value}' not in {self.device_mesh=}. " + "Set expert_parallel_degree > 1 in the device_mesh config." + ) + return self + + class DebuggingEnrichedModelConfig(BaseModel): model: PydanticPytorchModuleOrListType logging_dir_path: Path diff --git a/src/modalities/loss_functions.py b/src/modalities/loss_functions.py index e3be6100d..2412f4d59 100644 --- a/src/modalities/loss_functions.py +++ b/src/modalities/loss_functions.py @@ -1,11 +1,58 @@ from abc import ABC, abstractmethod -from typing import overload +from typing import Callable, overload import torch -from torch.nn import CrossEntropyLoss +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed.fsdp import FSDPModule as FSDP2 +from torch.utils.checkpoint import checkpoint from modalities.batch import InferenceResultBatch +# PyTorch's default ignore index for cross-entropy loss. Tokens with this label are +# excluded from both the loss value and the (valid-)token normalization. +IGNORE_INDEX = -100 + + +def clm_cross_entropy_loss(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + """Pure-tensor causal-LM cross-entropy with mean reduction over valid tokens. + + This is a free function (not a bound method) so it can be handed to + ``torch.compile`` as a clean, ``self``-free callable, mirroring TorchTitan's + module-level ``cross_entropy_loss`` compile target + (torchtitan/components/loss.py). Tokens labelled ``IGNORE_INDEX`` are ignored. + + Args: + logits (torch.Tensor): Unnormalized predictions of shape (..., vocab_size). + labels (torch.Tensor): Target token ids, broadcastable to ``logits[..., 0]``. + + Returns: + torch.Tensor: Scalar mean cross-entropy loss. + """ + # move labels to correct device to enable model parallelism + labels = labels.to(logits.device) + logits = logits.contiguous() + labels = labels.contiguous().long() + # Flatten the tokens. We compute here, the loss per token. + return F.cross_entropy( + logits.view(-1, logits.size(-1)), labels.view(-1), reduction="mean", ignore_index=IGNORE_INDEX + ) + + +def clm_cross_entropy_loss_sum(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + """Pure-tensor causal-LM cross-entropy with *sum* reduction over valid tokens. + + Used by the chunked loss: summing per-chunk contributions and dividing by the + total valid-token count reproduces the global mean of :func:`clm_cross_entropy_loss`, + while allowing each chunk to be computed (and freed) independently. + """ + labels = labels.to(logits.device) + logits = logits.contiguous() + labels = labels.contiguous().long() + return F.cross_entropy( + logits.view(-1, logits.size(-1)), labels.view(-1), reduction="sum", ignore_index=IGNORE_INDEX + ) + class Loss(ABC): def __init__(self, tag: str): @@ -23,14 +70,56 @@ def __call__(self, forward_batch: InferenceResultBatch) -> torch.Tensor: """ raise NotImplementedError + def compile(self, backend: str = "inductor") -> None: + """Compile the pure-tensor computation of this loss in place. + + Mirrors TorchTitan's ``BaseLoss._maybe_compile``: only the numeric + tensor-in/tensor-out core (``self.fn``) is compiled, never the batch/ + container unpacking. Subclasses with a compile-friendly core override this. + + Args: + backend (str): torch.compile backend. Defaults to "inductor". + """ + raise NotImplementedError(f"{type(self).__name__} does not support loss compilation.") + + +class LossFactory: + """Factory that applies training-time transformations to loss functions, + mirroring :class:`~modalities.models.model_factory.ModelFactory`.""" + + @staticmethod + def get_compiled_loss(loss: Loss, backend: str = "inductor") -> Loss: + """Compile the pure-tensor core of the given loss in place and return it. + + Follows the same in-place-mutate-and-return contract as + ``ModelFactory.get_compiled_model``. Composes with any ``Loss`` that + implements ``compile`` (e.g. wrapping a chunked loss compiles its CE core). + + Args: + loss (Loss): The loss whose numeric core should be compiled. + backend (str): torch.compile backend. Defaults to "inductor". + + Returns: + Loss: The same loss instance with its ``fn`` compiled. + """ + loss.compile(backend=backend) + return loss + class CLMCrossEntropyLoss(Loss): def __init__(self, target_key: str, prediction_key: str, tag: str = "CLMCrossEntropyLoss"): super().__init__(tag) self.target_key = target_key self.prediction_key = prediction_key - # Mean over the tokens in the local-batch (batch per rank) - self.loss_fun = CrossEntropyLoss(reduction="mean") + # Pure-tensor core. Swapped for a compiled variant by `compile`. + # Mean over the (valid) tokens in the local-batch (batch per rank). + self.fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = clm_cross_entropy_loss + + def compile(self, backend: str = "inductor") -> None: + # Compile only the tensor core, not the InferenceResultBatch unpacking. + # Note: unlike model/block compilation we do not pass fullgraph=True here, + # matching TorchTitan's loss compile (torchtitan/components/loss.py). + self.fn = torch.compile(self.fn, backend=backend) @overload def __call__(self, forward_batch: InferenceResultBatch) -> torch.Tensor: @@ -42,14 +131,7 @@ def __call__(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor def __call__(self, *args, **kwargs) -> torch.Tensor: labels, lm_logits = self._parse_arguments(args, kwargs) - - # move labels to correct device to enable model parallelism - labels = labels.to(lm_logits.device) - shift_logits = lm_logits.contiguous() - shift_labels = labels.contiguous().long() - # Flatten the tokens. We compute here, the loss per token. - loss = self.loss_fun(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) - return loss + return self.fn(lm_logits, labels) def _parse_arguments( self, @@ -87,6 +169,121 @@ def _parse_arguments( return labels, lm_logits +class ChunkedCLMCrossEntropyLoss(Loss): + """Memory-efficient causal-LM cross-entropy that never materializes the full + ``[batch, seq_len, vocab_size]`` logits tensor. + + Same goal as TorchTitan's ``ChunkedLossWrapper`` (torchtitan/components/loss.py): + the language-model head is moved *out* of the model's forward pass and applied + chunk-by-chunk inside the loss, so peak activation memory of the head + the + float32 up-cast inside cross-entropy is reduced by roughly ``num_chunks``. For a + 131k vocabulary this is the single largest activation in the model. + + Mechanism (adapted to modalities' idioms): + 1. The model runs with ``skip_lm_head=True`` and returns the post-norm + hidden states ``[batch, seq_len, n_embd]`` under ``prediction_key``. + 2. Hidden states and labels are split into ``num_chunks`` along the + sequence dimension. + 3. Each chunk is pushed through the (referenced) ``lm_head`` and cross-entropy + inside ``torch.utils.checkpoint``, so the chunk's logits are freed after + the forward and recomputed on demand during backward. Only one chunk's + logits are alive at any time. + 4. Per-chunk *sum*-reduced losses are accumulated and divided by the global + valid-token count, which is numerically equal to the mean reduction of + :class:`CLMCrossEntropyLoss`. + + Unlike TorchTitan, we rely on ``torch.utils.checkpoint`` (recomputing the head in + backward) rather than a manual per-chunk backward + custom autograd bridge. This + keeps the implementation torch-native and consistent with modalities' existing + activation-checkpointing approach; the trade-off is one extra ``lm_head`` forward + per chunk during backward. + + Note: + Under FSDP2 with ``reshard_after_forward=True`` the ``lm_head`` must be its + own FSDP unit so that the per-chunk (and recomputed) head calls trigger the + parameter all-gather. Tensor-parallel loss-parallel cross-entropy is not + handled here (modalities' plain CE is not loss-parallel either). + """ + + def __init__( + self, + model: nn.Module, + target_key: str, + prediction_key: str, + num_chunks: int = 8, + tag: str = "ChunkedCLMCrossEntropyLoss", + ): + """ + Args: + model (nn.Module): The (already wrapped) model that owns the ``lm_head``. + Passed BY_REFERENCE so this loss can borrow the head and switch the + model into ``skip_lm_head`` mode. Must expose ``lm_head`` and + ``set_skip_lm_head`` (see GPT2LLM). + target_key (str): Key of the label tensor in the batch targets. + prediction_key (str): Key under which the model stores the hidden states. + num_chunks (int): Number of sequence-dimension chunks. Defaults to 8. + tag (str): Loss tag. Defaults to "ChunkedCLMCrossEntropyLoss". + """ + super().__init__(tag) + if not isinstance(model, nn.Module) or not hasattr(model, "lm_head") or not hasattr(model, "set_skip_lm_head"): + raise ValueError( + "ChunkedCLMCrossEntropyLoss requires a single nn.Module exposing `lm_head` and " + "`set_skip_lm_head` (e.g. GPT2LLM). Pipeline-parallel model parts are not supported." + ) + self.target_key = target_key + self.prediction_key = prediction_key + self.num_chunks = num_chunks + self._lm_head = model.lm_head + # Move the head out of the model's forward; the head is applied here instead. + model.set_skip_lm_head(True) + # Pure-tensor core (sum reduction). Swapped for a compiled variant by `compile`. + self.fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = clm_cross_entropy_loss_sum + + def compile(self, backend: str = "inductor") -> None: + # Compile only the cross-entropy core; the lm_head is intentionally left + # uncompiled (matches TorchTitan's chunked loss). + self.fn = torch.compile(self.fn, backend=backend) + + def _chunk_loss(self, hidden_chunk: torch.Tensor, label_chunk: torch.Tensor) -> torch.Tensor: + # Runs inside checkpoint: the chunk logits produced here are not stored for + # backward but recomputed, so peak memory holds only one chunk of logits. + logits = self._lm_head(hidden_chunk) + return self.fn(logits, label_chunk) + + def __call__(self, forward_batch: InferenceResultBatch) -> torch.Tensor: + hidden_states = forward_batch.get_predictions(self.prediction_key) + labels = forward_batch.get_targets(self.target_key).to(hidden_states.device) + + # Normalize by the global valid-token count so the summed per-chunk losses + # equal the mean over valid tokens (clamped to avoid div-by-zero on a fully + # masked micro-batch). + num_valid_tokens = (labels != IGNORE_INDEX).sum().clamp(min=1) + + hidden_chunks = torch.chunk(hidden_states, self.num_chunks, dim=1) + label_chunks = torch.chunk(labels, self.num_chunks, dim=1) + + # When the lm_head is its own FSDP2 unit, keep its parameters unsharded across + # all chunk (and recompute) calls to avoid a fresh all-gather per chunk, then + # restore the default behaviour afterwards. Mirrors TorchTitan's + # ChunkedLossWrapper FSDP handling. No-op when the head is not an FSDPModule + # (single-device / DDP / head folded into the root FSDP unit). + head_is_fsdp_unit = isinstance(self._lm_head, FSDP2) + if head_is_fsdp_unit: + self._lm_head.set_reshard_after_forward(False) + + try: + total_loss = hidden_states.new_zeros(()) + for hidden_chunk, label_chunk in zip(hidden_chunks, label_chunks): + # use_reentrant=False is required for correct grads with non-tensor + # closure state and is the recommended checkpoint variant. + total_loss = total_loss + checkpoint(self._chunk_loss, hidden_chunk, label_chunk, use_reentrant=False) + finally: + if head_is_fsdp_unit: + self._lm_head.set_reshard_after_forward(True) + self._lm_head.reshard() + return total_loss / num_valid_tokens + + def nce_loss( embedding1: torch.Tensor, embedding2: torch.Tensor, device: torch.device, is_asymmetric: bool, temperature: float ) -> torch.Tensor: diff --git a/src/modalities/models/components/moe/experts.py b/src/modalities/models/components/moe/experts.py index 03798263a..72b8ad3e2 100644 --- a/src/modalities/models/components/moe/experts.py +++ b/src/modalities/models/components/moe/experts.py @@ -22,6 +22,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from torch.distributed.tensor import DTensor class ExpertsBackend(str, Enum): @@ -137,32 +138,58 @@ def _use_grouped_mm(self, x: torch.Tensor) -> bool: and x.dtype in (torch.bfloat16, torch.float16) ) + def _local_weights(self) -> tuple[torch.Tensor, torch.Tensor]: + """ + Returns the weights as plain tensors holding this rank's experts. + + Under expert parallelism ``w1``/``w2`` are DTensors sharded along the expert dimension, and + the matmul kernels operate on the local shard. Without expert parallelism this is a no-op. + + Returns: + tuple[torch.Tensor, torch.Tensor]: The local ``(w1, w2)``. + """ + w1, w2 = self.w1, self.w2 + return ( + w1.to_local() if isinstance(w1, DTensor) else w1, + w2.to_local() if isinstance(w2, DTensor) else w2, + ) + 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``. + tokens_per_expert (torch.Tensor): Token count per expert, of shape ``(num_local_experts,)``, + integer dtype. Must sum to ``num_routed``. Under expert parallelism these are the + counts for the experts this rank owns, not all ``num_experts``. Returns: torch.Tensor: Expert outputs of shape ``(num_routed, n_embd)``. """ + w1, w2 = self._local_weights() 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) + hidden = squared_relu(torch._grouped_mm(x_sorted, w1.transpose(-2, -1), offs=offsets)) + return torch._grouped_mm(hidden, w2.transpose(-2, -1), offs=offsets) - return self._forward_looped(x_sorted, tokens_per_expert) + return self._forward_looped(x_sorted, tokens_per_expert, w1, w2) - def _forward_looped(self, x_sorted: torch.Tensor, tokens_per_expert: torch.Tensor) -> torch.Tensor: + def _forward_looped( + self, + x_sorted: torch.Tensor, + tokens_per_expert: torch.Tensor, + w1: torch.Tensor, + w2: 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,)``. + tokens_per_expert (torch.Tensor): Token count per expert, of shape ``(num_local_experts,)``. + w1 (torch.Tensor): The local up-projection weights. + w2 (torch.Tensor): The local down-projection weights. Returns: torch.Tensor: Expert outputs of shape ``(num_routed, n_embd)``. @@ -176,7 +203,7 @@ def _forward_looped(self, x_sorted: torch.Tensor, tokens_per_expert: torch.Tenso 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) + hidden = squared_relu(chunk @ w1[expert_idx].transpose(0, 1)) + outputs[start:stop] = hidden @ w2[expert_idx].transpose(0, 1) start = stop return outputs diff --git a/src/modalities/models/gpt2/gpt2_model.py b/src/modalities/models/gpt2/gpt2_model.py index 993221e2c..1d544dcf6 100644 --- a/src/modalities/models/gpt2/gpt2_model.py +++ b/src/modalities/models/gpt2/gpt2_model.py @@ -1059,6 +1059,10 @@ def __init__( self.n_embd = n_embd self.n_layer = n_layer self.poe_type = poe_type + # When True, forward returns the post-norm hidden states instead of logits so + # that a memory-efficient loss (e.g. ChunkedCLMCrossEntropyLoss) can apply the + # lm_head chunk-by-chunk. Toggled via `set_skip_lm_head`. + self._skip_lm_head = False assert vocab_size is not None assert sequence_length is not None @@ -1121,6 +1125,17 @@ def __init__( self.transformer.lm_head.weight ) # https://paperswithcode.com/method/weight-tying + @property + def lm_head(self) -> nn.Module: + """The language-model head. Exposed so a memory-efficient loss can apply it + chunk-by-chunk (see ChunkedCLMCrossEntropyLoss).""" + return self.transformer.lm_head + + def set_skip_lm_head(self, skip: bool) -> None: + """Toggle whether forward returns post-norm hidden states (True) instead of + logits (False). Used together with a loss that owns the lm_head.""" + self._skip_lm_head = skip + @property def has_tied_word_embeddings(self) -> bool: # In pipeline parallelism a stage's transformer may not contain the wte/lm_head submodules @@ -1205,6 +1220,10 @@ def forward_impl(self, inputs: torch.Tensor) -> torch.Tensor: 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 + # When skipping the head, return the post-norm hidden states so the loss can + # apply the lm_head chunk-by-chunk (memory-efficient path). + if self._skip_lm_head: + return h h = self.transformer.lm_head(h) if hasattr(self.transformer, "lm_head") else h return h diff --git a/src/modalities/models/model_factory.py b/src/modalities/models/model_factory.py index 62933794d..1e6ae3e72 100644 --- a/src/modalities/models/model_factory.py +++ b/src/modalities/models/model_factory.py @@ -31,6 +31,7 @@ from modalities.checkpointing.checkpoint_loading import FSDP1CheckpointLoadingIF from modalities.config.config import ActivationCheckpointedModelConfig from modalities.exceptions import ModelStateError +from modalities.models.components.moe.moe import MoE from modalities.models.gpt2.gpt2_model import ( GPT2LLM, AttentionConfig, @@ -41,9 +42,14 @@ TransformerMLP, ) from modalities.models.model import ActivationType +from modalities.models.parallelism.expert_parallelism import ExpertParallelGroupedExperts, shard_experts_over_ep_mesh from modalities.nn.model_initialization.initialization_if import ModelInitializationIF from modalities.running_env.env_utils import FSDP2MixedPrecisionSettings, MixedPrecisionSettings -from modalities.running_env.fsdp.device_mesh import ParallelismDegrees +from modalities.running_env.fsdp.device_mesh import ( + ParallelismDegrees, + get_mesh_for_parallelism_method, + has_parallelism_method, +) from modalities.running_env.fsdp.fsdp_auto_wrapper import FSDPTransformerAutoWrapPolicyFactory from modalities.training.activation_checkpointing.activation_checkpointing import ( ActivationCheckpointing, @@ -165,6 +171,53 @@ def get_fsdp1_wrapped_model( ) return fsdp_model + @staticmethod + def get_expert_parallelized_model(model: nn.Module, device_mesh: DeviceMesh) -> nn.Module: + """Applies expert parallelism to every mixture-of-experts layer of the model, in place. + + Each MoE layer's routed expert stack is partitioned along the expert dimension over the + ``ep`` mesh dimension and wrapped so that tokens are dispatched to, and combined from, the + rank owning their expert. Shared experts, the router and all dense layers are untouched. + + Must run *before* FSDP2 wrapping (which then shards what is left of the data-parallel + dimension on top) and while the model is still on the meta device, since the expert + parameters are re-created rather than redistributed. + + Args: + model (nn.Module): The model to parallelize, expected to be on the meta device. + device_mesh (DeviceMesh): The device mesh. Must have an ``ep`` dimension, i.e. the + device mesh config must set ``expert_parallel_degree > 1``. + + Raises: + ModelStateError: If the mesh has no ``ep`` dimension, or if the model has no MoE layer. + + Returns: + nn.Module: The same model, with its expert stacks replaced in place. + """ + if not has_parallelism_method(device_mesh, ParallelismDegrees.EP): + raise ModelStateError( + "Expert parallelism requires an `ep` dimension in the device mesh. Set " + "expert_parallel_degree > 1 in the device_mesh config." + ) + ep_mesh = get_mesh_for_parallelism_method(device_mesh, ParallelismDegrees.EP) + + moe_layers = [module for module in model.modules() if isinstance(module, MoE)] + if not moe_layers: + raise ModelStateError( + "Expert parallelism was requested but the model contains no MoE layer. Remove the " + "expert-parallelized model component or set expert_parallel_degree to 1." + ) + for moe in moe_layers: + if isinstance(moe.experts, ExpertParallelGroupedExperts): + continue + shard_experts_over_ep_mesh(moe.experts, ep_mesh) + moe.experts = ExpertParallelGroupedExperts(experts=moe.experts, ep_mesh=ep_mesh) + logger.info( + f"Applied expert parallelism (degree {ep_mesh.size()}) to {len(moe_layers)} MoE layer(s); " + f"{moe_layers[0].router.num_experts // ep_mesh.size()} local experts per rank." + ) + return model + @staticmethod def get_fsdp2_wrapped_model( model: nn.Module, @@ -173,6 +226,7 @@ def get_fsdp2_wrapped_model( mixed_precision_settings: FSDP2MixedPrecisionSettings, reshard_after_forward: bool, layers_per_fsdp_unit: int = 1, + separate_lm_head_fsdp_unit: bool = False, ) -> FSDP2: """Get the FSDP2-wrapped model. @@ -212,6 +266,28 @@ def get_fsdp2_wrapped_model( modules = list(model.modules()) + # Expert-parallel stacks must not be sharded on the full dp_shard dimension: ranks differing + # in their `ep` coordinate hold *different* experts, so all-gathering across them would mix + # unrelated weights. Give each expert stack its own FSDP unit on dp_shard_mod_ep, the part of + # the data-parallel dimension that expert parallelism did not consume. Nested FSDP units are + # skipped by the enclosing ones, so the block-level sharding below leaves them alone. + # At expert_parallel_degree == data_parallel_shard_degree this mesh has size 1 and the unit + # only exists to claim the parameters; there is nothing left to shard or reduce. + expert_modules = [m for m in modules if isinstance(m, ExpertParallelGroupedExperts)] + if expert_modules: + expert_fsdp_mesh = get_mesh_for_parallelism_method(device_mesh, ParallelismDegrees.DP_SHARD_MOD_EP) + for expert_module in expert_modules: + fully_shard( + expert_module.experts, + mesh=expert_fsdp_mesh, + mp_policy=mp_policy, + reshard_after_forward=reshard_after_forward, + ) + logger.info( + f"Sharded {len(expert_modules)} expert-parallel stack(s) on " + f"dp_shard_mod_ep (size {expert_fsdp_mesh.size()})." + ) + # we first shard all the blocks grouped_modules: list[nn.Module] = [] module_id = 0 @@ -237,6 +313,16 @@ def get_fsdp2_wrapped_model( reshard_after_forward=reshard_block_after_forward, ) + # Optionally shard the lm_head as its own FSDP unit. Required by the + # memory-efficient chunked loss (ChunkedCLMCrossEntropyLoss): it applies the + # head outside the model's forward, so the head needs its own gather hook to + # all-gather its parameters when called (and recomputed) per chunk. Without + # this, a resharded head would silently run on a local shard. + if separate_lm_head_fsdp_unit: + if not hasattr(model, "lm_head"): + raise ModelStateError("separate_lm_head_fsdp_unit=True requires the model to expose `lm_head`.") + fully_shard(model.lm_head, **fsdp_config, reshard_after_forward=reshard_after_forward) + # finally, we shard the entire model fully_shard(model, **fsdp_config, reshard_after_forward=reshard_after_forward) logger.info( diff --git a/src/modalities/models/parallelism/expert_parallelism.py b/src/modalities/models/parallelism/expert_parallelism.py new file mode 100644 index 000000000..75c5f3b2b --- /dev/null +++ b/src/modalities/models/parallelism/expert_parallelism.py @@ -0,0 +1,287 @@ +"""Expert parallelism for mixture-of-experts layers. + +Under expert parallelism the routed experts of a layer are partitioned across the ``ep`` mesh +dimension, so each rank stores and evaluates only ``num_experts / ep_degree`` of them. Because every +rank still routes its *own* tokens, tokens have to travel to the rank that owns their expert and +their results have to travel back. That is the dispatch/combine pair of all-to-all collectives this +module inserts around :class:`~modalities.models.components.moe.experts.GroupedExperts`. + +The alternative -- what modalities does without expert parallelism -- is to let FSDP2 all-gather the +full expert stack on every rank. For Nemotron-3 Nano 30B-A3B that means moving 2.55 GiB of bf16 +expert weights per MoE layer per forward, versus roughly a sixth of that in tokens here, which is +what makes expert parallelism worthwhile at scale. + +Design notes +------------ +The dispatch/combine structure follows Meta's open-source project TorchTitan +(``torchtitan/models/common/moe.py``, ``torchtitan/distributed/expert_parallel.py``), licensed under +the BSD 3-Clause License. + +Two things are done deliberately differently from a straightforward implementation: + +1. **The post-all-to-all permutation is built entirely on device.** Received tokens arrive grouped by + *sender* and have to be regrouped by *local expert* before a grouped matmul can consume them. The + obvious implementation loops over ``(sender, local_expert)`` pairs and calls ``.item()`` to slice + each range, which costs ``2 * ep_degree * num_local_experts`` device-to-host synchronizations per + MoE layer per forward pass -- several thousand per step for a 23-MoE-layer model. Here the index + tensor is constructed with vectorized ops instead (see :func:`_build_permute_indices`). +2. **There is exactly one device-to-host synchronization per dispatch**, for the split lists that + ``all_to_all_single`` requires on the host. Both split vectors are copied in a single transfer. + +Autograd flows through both all-to-all calls via ``all_to_all_single_autograd`` (the adjoint of an +all-to-all is an all-to-all with the splits swapped) and through the gather/scatter permutation. +""" + +import torch +import torch.nn as nn +from torch.distributed._functional_collectives import all_to_all_single, all_to_all_single_autograd +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor, Shard + +from modalities.models.components.moe.experts import GroupedExperts + + +def _build_permute_indices( + tokens_per_expert_group: torch.Tensor, + ep_degree: int, + num_local_experts: int, + total_tokens: int, +) -> torch.Tensor: + """ + Builds the index tensor that regroups received tokens from sender-major to expert-major order. + + After the token all-to-all, the receive buffer is ordered by sender and, within a sender, by + local expert:: + + [r0e0, r0e1, ..., r0e(L-1)] [r1e0, r1e1, ...] ... [r(P-1)e0, ...] + + A grouped matmul needs all tokens of one expert to be contiguous, i.e. expert-major order:: + + [e0r0, e0r1, ..., e0r(P-1)] [e1r0, ...] ... [e(L-1)r0, ...] + + Both orders are concatenations of the same ``P * L`` variable-length blocks, so the permutation + is fully determined by the block lengths. This function turns those lengths into the flat index + tensor without a single host synchronization. + + Args: + tokens_per_expert_group (torch.Tensor): Received token counts of shape + ``(ep_degree * num_local_experts,)``, ordered by sender then by local expert. + ep_degree (int): The expert parallel degree, i.e. the number of senders ``P``. + num_local_experts (int): The number of experts owned by this rank, ``L``. + total_tokens (int): The sum of ``tokens_per_expert_group``. Passed in as a Python int because + the caller already has it on the host from the all-to-all split lists; recomputing it + here would reintroduce a synchronization. + + Returns: + torch.Tensor: Int64 gather indices of shape ``(total_tokens,)`` such that + ``received[indices]`` is in expert-major order. + """ + device = tokens_per_expert_group.device + counts = tokens_per_expert_group.view(ep_degree, num_local_experts) + # Exclusive cumulative sum over the receive buffer gives each block's start offset. + flat_counts = counts.reshape(-1) + block_starts = flat_counts.cumsum(0) - flat_counts + + # Transpose both into expert-major (expert, sender) order, which is the output order. + counts_out = counts.t().reshape(-1) + starts_out = block_starts.view(ep_degree, num_local_experts).t().reshape(-1) + + if total_tokens == 0: + return torch.zeros(0, dtype=torch.long, device=device) + + # For every output position, find which block it belongs to and its offset within that block. + # repeat_interleave needs output_size to stay synchronization-free. + block_of_position = torch.repeat_interleave( + torch.arange(counts_out.numel(), device=device), + counts_out, + output_size=total_tokens, + ) + offset_in_block = torch.arange(total_tokens, device=device) - (counts_out.cumsum(0) - counts_out)[block_of_position] + return starts_out[block_of_position] + offset_in_block + + +class ExpertParallelGroupedExperts(nn.Module): + """ + Wraps a :class:`GroupedExperts` stack whose experts are partitioned across the ``ep`` mesh dim. + + Drop-in replacement for the wrapped module: the forward signature + ``(tokens_sorted_by_expert, tokens_per_expert)`` and the returned shape are unchanged, so + :class:`~modalities.models.components.moe.moe.MoE` needs no knowledge of expert parallelism. The + ``tokens_per_expert`` it passes in are *global* counts over all ``num_experts``; the wrapped + module receives the local counts for the experts this rank owns. + """ + + def __init__(self, experts: GroupedExperts, ep_mesh: DeviceMesh): + """ + Initializes the expert-parallel wrapper. + + Args: + experts (GroupedExperts): The expert stack. Its ``w1``/``w2`` are expected to already be + sharded over ``ep_mesh`` along the expert dimension (see + :func:`shard_experts_over_ep_mesh`). + ep_mesh (DeviceMesh): The one-dimensional ``ep`` sub-mesh whose process group carries the + dispatch and combine all-to-all collectives. + + Raises: + ValueError: If the mesh is not one-dimensional, or if the number of experts is not + divisible by the expert parallel degree. + """ + super().__init__() + if ep_mesh.ndim != 1: + raise ValueError(f"ep_mesh must be one-dimensional, got ndim={ep_mesh.ndim}.") + ep_degree = ep_mesh.size() + if experts.num_experts % ep_degree != 0: + raise ValueError( + f"num_experts ({experts.num_experts}) must be divisible by the expert parallel " + f"degree ({ep_degree})." + ) + + self.experts = experts + self.ep_mesh = ep_mesh + self.ep_degree = ep_degree + self.num_local_experts = experts.num_experts // ep_degree + + @property + def num_experts(self) -> int: + """The *global* number of experts, which is what the router is validated against.""" + return self.experts.num_experts + + def _dispatch( + self, x_sorted: torch.Tensor, tokens_per_expert: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[int], list[int]]: + """ + Sends every token to the rank owning its expert and regroups the receive buffer by expert. + + Args: + x_sorted (torch.Tensor): Locally routed tokens sorted by global expert index, of shape + ``(num_routed, n_embd)``. + tokens_per_expert (torch.Tensor): Local token counts per *global* expert, of shape + ``(num_experts,)``. + + Returns: + tuple: ``(x_received, local_tokens_per_expert, permute_indices, input_splits, + output_splits)``, where ``x_received`` is in expert-major order and + ``local_tokens_per_expert`` has shape ``(num_local_experts,)``. + """ + group = self.ep_mesh.get_group() + + with torch.no_grad(): + # Experts are assigned to ranks in contiguous blocks, matching Shard(0) on the expert + # dimension, so an even-split all-to-all sends each rank exactly the counts of the + # experts it owns. + tokens_per_expert_group = all_to_all_single(tokens_per_expert, None, None, group=group) + tokens_per_expert_group = torch.ops._c10d_functional.wait_tensor(tokens_per_expert_group) + # all_to_all_single takes its split lists on the host. Stacking both vectors keeps this + # to a single device-to-host copy, which is the only synchronization in the dispatch. + splits = torch.stack( + ( + tokens_per_expert.view(self.ep_degree, -1).sum(dim=1), + tokens_per_expert_group.view(self.ep_degree, -1).sum(dim=1), + ) + ).to(device="cpu", non_blocking=False) + input_splits: list[int] = splits[0].tolist() + output_splits: list[int] = splits[1].tolist() + + x_received = all_to_all_single_autograd(x_sorted, output_splits, input_splits, group) + permute_indices = _build_permute_indices( + tokens_per_expert_group=tokens_per_expert_group, + ep_degree=self.ep_degree, + num_local_experts=self.num_local_experts, + total_tokens=sum(output_splits), + ) + local_tokens_per_expert = tokens_per_expert_group.view(self.ep_degree, self.num_local_experts).sum(dim=0) + return ( + x_received.index_select(0, permute_indices), + local_tokens_per_expert, + permute_indices, + input_splits, + output_splits, + ) + + def _combine( + self, + expert_out: torch.Tensor, + permute_indices: torch.Tensor, + num_received: int, + input_splits: list[int], + output_splits: list[int], + ) -> torch.Tensor: + """ + Undoes the expert-major regrouping and returns each token's result to its owning rank. + + Args: + expert_out (torch.Tensor): Expert outputs in expert-major order. + permute_indices (torch.Tensor): The indices produced by :meth:`_dispatch`. + num_received (int): Row count of the dispatch receive buffer. + input_splits (list[int]): The dispatch input splits; the combine output splits. + output_splits (list[int]): The dispatch output splits; the combine input splits. + + Returns: + torch.Tensor: Results in the caller's original expert-sorted order. + """ + # Scatter back to sender-major order. index_copy is the differentiable inverse of the + # index_select used in the dispatch. + unpermuted = expert_out.new_zeros((num_received, expert_out.shape[-1])).index_copy( + 0, permute_indices, expert_out + ) + # Splits swap relative to the dispatch: what we received, we now send back. + return all_to_all_single_autograd(unpermuted, input_splits, output_splits, self.ep_mesh.get_group()) + + def forward(self, x_sorted: torch.Tensor, tokens_per_expert: torch.Tensor) -> torch.Tensor: + """ + Evaluates the globally distributed expert stack on locally routed tokens. + + Args: + x_sorted (torch.Tensor): Tokens sorted by global expert index, shape + ``(num_routed, n_embd)``. + tokens_per_expert (torch.Tensor): Token counts per global expert, shape + ``(num_experts,)``. + + Returns: + torch.Tensor: Expert outputs of shape ``(num_routed, n_embd)``, in the same order as + ``x_sorted``. + """ + x_received, local_tokens_per_expert, permute_indices, input_splits, output_splits = self._dispatch( + x_sorted, tokens_per_expert + ) + expert_out = self.experts(x_received, local_tokens_per_expert) + return self._combine( + expert_out=expert_out, + permute_indices=permute_indices, + num_received=x_received.shape[0], + input_splits=input_splits, + output_splits=output_splits, + ) + + +def shard_experts_over_ep_mesh(experts: GroupedExperts, ep_mesh: DeviceMesh) -> None: + """ + Replaces an expert stack's weights in place with DTensors sharded over the expert dimension. + + The parameters are re-created rather than redistributed: this runs on a meta-device model whose + weights are materialized and initialized later, so there is no data to preserve. Keeping them as + DTensors (rather than plain per-rank slices) means checkpoints still see the full global expert + stack, and lets FSDP2 shard what remains of the data-parallel dimension on top. + + Args: + experts (GroupedExperts): The expert stack to shard in place. + ep_mesh (DeviceMesh): The one-dimensional ``ep`` sub-mesh to shard over. + + Raises: + ValueError: If a weight's expert dimension is not divisible by the expert parallel degree. + """ + ep_degree = ep_mesh.size() + for name in ("w1", "w2"): + param = getattr(experts, name) + if isinstance(param, DTensor): + continue + num_experts = param.shape[0] + if num_experts % ep_degree != 0: + raise ValueError( + f"Expert dimension of {name} ({num_experts}) is not divisible by the expert " + f"parallel degree ({ep_degree})." + ) + local_shape = (num_experts // ep_degree, *param.shape[1:]) + local = torch.empty(local_shape, dtype=param.dtype, device=param.device) + sharded = DTensor.from_local(local, ep_mesh, [Shard(0)], run_check=False) + experts.register_parameter(name, nn.Parameter(sharded, requires_grad=param.requires_grad)) diff --git a/src/modalities/registry/components.py b/src/modalities/registry/components.py index 6bc548d93..db43b3d32 100644 --- a/src/modalities/registry/components.py +++ b/src/modalities/registry/components.py @@ -24,8 +24,10 @@ AdamWOptimizerConfig, BatchSamplerConfig, CheckpointSavingConfig, + ChunkedCLMCrossEntropyLossConfig, CLMCrossEntropyLossConfig, CombinedDatasetConfig, + CompiledLossConfig, CompiledModelConfig, ConstantLRSchedulerConfig, CosineAnnealingLRSchedulerConfig, @@ -37,6 +39,7 @@ DummyProgressSubscriberConfig, DummyResultSubscriberConfig, EvaluationResultToDiscSubscriberConfig, + ExpertParallelizedModelConfig, FSDP1ActivationCheckpointedModelConfig, FSDP1CheckpointedModelConfig, FSDP1CheckpointedOptimizerConfig, @@ -84,7 +87,7 @@ ProgressSubscriberFactory, ResultsSubscriberFactory, ) -from modalities.loss_functions import CLMCrossEntropyLoss +from modalities.loss_functions import ChunkedCLMCrossEntropyLoss, CLMCrossEntropyLoss, LossFactory from modalities.models.coca.coca_model import CoCa, CoCaConfig from modalities.models.coca.collator import CoCaCollateFnConfig, CoCaCollatorFn from modalities.models.components.layer_norms import ( @@ -223,6 +226,12 @@ class ComponentEntity: ComponentEntity( "model", "fsdp2_wrapped", maybe_model_list(ModelFactory.get_fsdp2_wrapped_model), FSDP2WrappedModelConfig ), + ComponentEntity( + "model", + "expert_parallelized", + maybe_model_list(ModelFactory.get_expert_parallelized_model), + ExpertParallelizedModelConfig, + ), ComponentEntity( "model", "model_initialized", @@ -284,6 +293,10 @@ class ComponentEntity: ), # losses ComponentEntity("loss", "clm_cross_entropy_loss", CLMCrossEntropyLoss, CLMCrossEntropyLossConfig), + ComponentEntity( + "loss", "chunked_clm_cross_entropy_loss", ChunkedCLMCrossEntropyLoss, ChunkedCLMCrossEntropyLossConfig + ), + ComponentEntity("loss", "compiled", LossFactory.get_compiled_loss, CompiledLossConfig), ComponentEntity("loss", "moe_aux_loss", MoEAuxLoss, MoEAuxLossConfig), ComponentEntity("loss", "weighted_sum", WeightedSumLoss, WeightedSumLossConfig), # optimizers diff --git a/src/modalities/running_env/fsdp/device_mesh.py b/src/modalities/running_env/fsdp/device_mesh.py index cd456938c..9b7c992f9 100644 --- a/src/modalities/running_env/fsdp/device_mesh.py +++ b/src/modalities/running_env/fsdp/device_mesh.py @@ -21,6 +21,12 @@ class DeviceMeshConfig(BaseModel): tensor_parallel_degree: Annotated[int, Field(strict=True, gt=0)] = 1 pipeline_parallel_degree: Annotated[int, Field(strict=True, gt=0)] = 1 context_parallel_degree: Annotated[int, Field(strict=True, gt=0)] = 1 + # Expert parallelism is carved *out of* the data-parallel shard dimension rather than + # multiplying into the world size: EP ranks each hold a distinct slice of the experts but + # still process their own data shard, exchanging tokens via all-to-all. The number of + # distinct data shards therefore stays data_parallel_shard_degree * data_parallel_replicate_degree, + # which is what the dataloader's dp_degree must keep seeing. + expert_parallel_degree: Annotated[int, Field(strict=True, gt=0)] = 1 enable_loss_parallel: Optional[bool] = False world_size: Annotated[int, Field(strict=True, gt=0)] @@ -78,12 +84,38 @@ def _validate(self): ) if self.enable_loss_parallel and self.tensor_parallel_degree <= 1: raise ConfigError(f"{self.enable_loss_parallel=} requires tensor_parallel_degree > 1") + + if self.expert_parallel_degree > 1: + if self.expert_parallel_degree > self.data_parallel_shard_degree: + raise ConfigError( + f"expert_parallel_degree({self.expert_parallel_degree}) must not exceed " + f"data_parallel_shard_degree({self.data_parallel_shard_degree}): expert parallelism " + "is carved out of the data-parallel shard dimension." + ) + if self.data_parallel_shard_degree % self.expert_parallel_degree != 0: + raise ConfigError( + f"data_parallel_shard_degree({self.data_parallel_shard_degree}) must be divisible by " + f"expert_parallel_degree({self.expert_parallel_degree})." + ) + for degree, name in ( + (self.tensor_parallel_degree, "tensor_parallel_degree"), + (self.pipeline_parallel_degree, "pipeline_parallel_degree"), + (self.context_parallel_degree, "context_parallel_degree"), + ): + if degree > 1: + raise ConfigError(f"expert_parallel_degree > 1 is not yet supported together with {name}={degree}.") return self class ParallelismDegrees(Enum): DP_REPLICATE = "dp_replicate" DP_SHARD = "dp_shard" + # Only present when expert_parallel_degree > 1, where DP_SHARD is split into these two + # sub-dimensions and re-exposed as a flattened DP_SHARD alias. EP is the inner (fastest + # varying) of the two so that an all-to-all group spans consecutive global ranks, i.e. stays + # inside a node for the common case of expert_parallel_degree <= devices per node. + DP_SHARD_MOD_EP = "dp_shard_mod_ep" + EP = "ep" CP = "cp" TP = "tp" PP = "pp" @@ -98,6 +130,7 @@ def get_device_mesh( context_parallel_degree: int, enable_loss_parallel: bool, world_size: int, + expert_parallel_degree: int = 1, ) -> DeviceMesh: """ Gets the device mesh for the specified parallelism degrees. @@ -111,40 +144,75 @@ def get_device_mesh( context_parallel_degree (int): The context parallel degree. enable_loss_parallel (bool): Whether to enable loss parallelism. world_size (int): The world size. + expert_parallel_degree (int): The expert parallel degree. Carved out of + ``data_parallel_shard_degree`` (see :class:`DeviceMeshConfig`), so a degree > 1 splits + the ``dp_shard`` dimension into ``dp_shard_mod_ep`` x ``ep``. The original ``dp_shard`` + stays available as a flattened dimension, so callers that only care about data + parallelism need no changes. Returns: DeviceMesh: The device mesh. """ + # With expert parallelism, dp_shard is materialized as two dimensions and flattened back + # afterwards; without it, the mesh is built exactly as before. + if expert_parallel_degree > 1: + dp_shard_dims = [data_parallel_shard_degree // expert_parallel_degree, expert_parallel_degree] + dp_shard_names = [ParallelismDegrees.DP_SHARD_MOD_EP.value, ParallelismDegrees.EP.value] + else: + dp_shard_dims = [data_parallel_shard_degree] + dp_shard_names = [ParallelismDegrees.DP_SHARD.value] + dims = [] names = [] for dim, name in zip( - [ - pipeline_parallel_degree, - data_parallel_replicate_degree, - data_parallel_shard_degree, - context_parallel_degree, - tensor_parallel_degree, - ], - [ - ParallelismDegrees.PP.value, - ParallelismDegrees.DP_REPLICATE.value, - ParallelismDegrees.DP_SHARD.value, - ParallelismDegrees.CP.value, - ParallelismDegrees.TP.value, - ], + [pipeline_parallel_degree, data_parallel_replicate_degree] + + dp_shard_dims + + [context_parallel_degree, tensor_parallel_degree], + [ParallelismDegrees.PP.value, ParallelismDegrees.DP_REPLICATE.value] + + dp_shard_names + + [ParallelismDegrees.CP.value, ParallelismDegrees.TP.value], strict=True, ): - if dim > 1 or name == ParallelismDegrees.DP_SHARD.value: + # The dp_shard sub-dimensions are always kept, even at degree 1, because the flattening + # below has to address both of them by name. + if dim > 1 or name in dp_shard_names: dims.append(dim) names.append(name) names = tuple(names) device_mesh = init_device_mesh(device_type, dims, mesh_dim_names=names) - logger.info(f"{device_mesh=} | {world_size=} | {enable_loss_parallel=}") + if expert_parallel_degree > 1: + # Re-expose the full data-parallel shard group under its original name. This is a lookup + # alias (device_mesh["dp_shard"]); it deliberately does not appear in mesh_dim_names, which + # is why the helpers below resolve names through __getitem__ instead. + device_mesh[tuple(dp_shard_names)]._flatten(ParallelismDegrees.DP_SHARD.value) + logger.info(f"{device_mesh=} | {world_size=} | {enable_loss_parallel=} | {expert_parallel_degree=}") # TODO: Torch Titan had some more checks here. We need to check if we also need those: # https://github.com/pytorch/torchtitan/blob/b291ad662493b63d25b038a30a915082d3617baf/torchtitan/distributed/parallel_dims.py#L86-L104 return device_mesh +def _resolve_sub_mesh(device_mesh: DeviceMesh | None, parallelism_method: ParallelismDegrees) -> DeviceMesh | None: + """Resolves a mesh dimension by name, returning None if the mesh does not have it. + + Resolution goes through ``device_mesh[name]`` rather than ``mesh_dim_names`` so that flattened + dimensions are found too. Under expert parallelism ``dp_shard`` is such a flattened alias: it is + addressable but absent from ``mesh_dim_names``. + + Args: + device_mesh (DeviceMesh | None): The device mesh. + parallelism_method (ParallelismDegrees): The parallelism method to resolve. + + Returns: + DeviceMesh | None: The sub-mesh, or None if this mesh has no such dimension. + """ + if device_mesh is None or device_mesh.mesh_dim_names is None: + return None + try: + return device_mesh[parallelism_method.value] + except (KeyError, RuntimeError, IndexError): + return None + + def get_parallel_degree(device_mesh: DeviceMesh, parallelism_methods: list[ParallelismDegrees]) -> int: """Gets the number of parallel ranks (i.e., the parallelism degree) from the device mesh for a specific parallelism method. @@ -158,9 +226,9 @@ def get_parallel_degree(device_mesh: DeviceMesh, parallelism_methods: list[Paral raise ValueError("device_mesh.mesh_dim_names is None") return prod( - device_mesh.size(device_mesh.mesh_dim_names.index(method.value)) + sub_mesh.size() for method in parallelism_methods - if method.value in device_mesh.mesh_dim_names + if (sub_mesh := _resolve_sub_mesh(device_mesh, method)) is not None ) @@ -174,11 +242,7 @@ def has_parallelism_method(device_mesh: DeviceMesh | None, parallelism_method: P Returns: bool: True if the device mesh has the specified parallelism method, False otherwise. """ - return ( - device_mesh is not None - and (mesh_dim_names := device_mesh.mesh_dim_names) is not None - and parallelism_method.value in mesh_dim_names - ) + return _resolve_sub_mesh(device_mesh, parallelism_method) is not None def get_mesh_for_parallelism_method(device_mesh: DeviceMesh, parallelism_method: ParallelismDegrees) -> DeviceMesh: diff --git a/src/modalities/training/gradient_clipping/fsdp_gradient_clipper.py b/src/modalities/training/gradient_clipping/fsdp_gradient_clipper.py index a5c0d2cfe..6ddfdc75d 100644 --- a/src/modalities/training/gradient_clipping/fsdp_gradient_clipper.py +++ b/src/modalities/training/gradient_clipping/fsdp_gradient_clipper.py @@ -17,6 +17,30 @@ from modalities.training.gradient_clipping.gradient_clipper import GradientClipperIF +def _group_by_gradient_mesh(parameters: list[torch.nn.Parameter]) -> list[list[torch.nn.Parameter]]: + """Groups parameters by the device mesh their gradient lives on. + + Both the norm reduction and the in-place rescaling of gradients are batched (``torch.stack`` and + ``aten._foreach_mul_`` respectively), and neither has a sharding rule for operands from different + meshes. Under expert parallelism the routed expert gradients are sharded over + ``(dp_shard_mod_ep, ep)`` while all other gradients are sharded over ``dp_shard``, so the batched + calls have to be made per mesh. Without expert parallelism this returns a single group and every + caller behaves exactly as before. + + Args: + parameters (list[torch.nn.Parameter]): Parameters with a gradient. + + Returns: + list[list[torch.nn.Parameter]]: One list of parameters per distinct gradient mesh. + """ + groups: dict[DeviceMesh | None, list[torch.nn.Parameter]] = {} + for parameter in parameters: + # Meshes are hashable and identity-stable, so they can key the grouping directly. + key = parameter.grad.device_mesh if isinstance(parameter.grad, DTensor) else None + groups.setdefault(key, []).append(parameter) + return list(groups.values()) + + class GradientClippingMode(LookupEnum): """ Enum class representing different modes of gradient clipping. @@ -132,6 +156,46 @@ def __init__( self.error_if_nonfinite = error_if_nonfinite self.foreach = foreach + def _get_total_norm_across_meshes(self, parameter_groups: list[list[torch.nn.Parameter]]) -> torch.Tensor: + """Computes one global gradient norm over gradients that may live on different device meshes. + + Norms decompose over any partition of the tensors, so each group (see + :func:`_group_by_gradient_mesh`) is reduced to a scalar on its own mesh and the group norms + are then combined: ``(sum_g norm_g ** p) ** (1/p)`` for a p-norm and ``max_g norm_g`` for the + infinity norm. With a single group this is exactly the previous computation. + + Args: + parameter_groups (list[list[torch.nn.Parameter]]): Parameters grouped by gradient mesh. + + Returns: + torch.Tensor: The global gradient norm as a plain (non-DTensor) scalar tensor. + """ + group_norms: list[torch.Tensor] = [] + for parameters in parameter_groups: + group_norm = torch.nn.utils.get_total_norm( + tensors=[parameter.grad for parameter in parameters], + norm_type=self.norm_type.value, + error_if_nonfinite=self.error_if_nonfinite, + foreach=self.foreach, + ) + # Inspired by torch titan + # If the norm is a DTensor, the placements must be + # `torch.distributed._tensor.ops.math_ops._NormPartial`. Reducing the DTensor yields the + # total norm over this group's process groups; converting to a local tensor then gives a + # value whose .item() is correct. + if isinstance(group_norm, DTensor): + # Will reach here if any non-PP parallelism is used. + # If only using PP, the norm will be a local tensor. + group_norm = group_norm.full_tensor() + group_norms.append(group_norm) + + if len(group_norms) == 1: + return group_norms[0] + stacked = torch.stack(group_norms) + if math.isinf(self.norm_type.value): + return stacked.max() + return stacked.pow(self.norm_type.value).sum().pow(1.0 / self.norm_type.value) + @torch.no_grad() def clip_gradients(self) -> torch.Tensor: """ @@ -140,23 +204,8 @@ def clip_gradients(self) -> torch.Tensor: Returns: torch.Tensor: The gradient norms. """ - grads = [p.grad for model in self.models for p in model.parameters() if p.grad is not None] - total_norm = torch.nn.utils.get_total_norm( - tensors=grads, - norm_type=self.norm_type.value, - error_if_nonfinite=self.error_if_nonfinite, - foreach=self.foreach, - ) - - # Inspired by torch titan - # If total_norm is a DTensor, the placements must be `torch.distributed._tensor.ops.math_ops._NormPartial`. - # We can simply reduce the DTensor to get the total norm in this tensor's process group - # and then convert it to a local tensor. - # NOTE: It has the purpose to return a reduced total_norm tensor whose .item() would return the correct value - if isinstance(total_norm, DTensor): - # Will reach here if any non-PP parallelism is used. - # If only using PP, total_norm will be a local tensor. - total_norm = total_norm.full_tensor() + parameters = [p for model in self.models for p in model.parameters() if p.grad is not None] + total_norm = self._get_total_norm_across_meshes(_group_by_gradient_mesh(parameters)) if has_parallelism_method(self.device_mesh, ParallelismDegrees.PP): pp_mesh = get_mesh_for_parallelism_method( @@ -221,10 +270,14 @@ def clip_gradients(self) -> torch.Tensor: """ total_norm = super().clip_gradients() for model in self.models: - torch.nn.utils.clip_grads_with_norm_( - parameters=model.parameters(), - max_norm=self.max_norm, - total_norm=total_norm, - foreach=self.foreach, - ) + parameters = [p for p in model.parameters() if p.grad is not None] + # Rescaling is batched per mesh: clip_grads_with_norm_ ends in a single + # aten._foreach_mul_ over all gradients, which cannot mix meshes. + for parameter_group in _group_by_gradient_mesh(parameters): + torch.nn.utils.clip_grads_with_norm_( + parameters=parameter_group, + max_norm=self.max_norm, + total_norm=total_norm, + foreach=self.foreach, + ) return total_norm diff --git a/tests/config/test_device_mesh_config.py b/tests/config/test_device_mesh_config.py new file mode 100644 index 000000000..f71c1a0dd --- /dev/null +++ b/tests/config/test_device_mesh_config.py @@ -0,0 +1,50 @@ +"""Validation tests for :class:`DeviceMeshConfig`, focused on the expert parallel degree. + +Expert parallelism is carved out of the data-parallel shard dimension rather than multiplying into +the world size, which makes its constraints easy to get wrong in a config. These tests pin them down +without needing a process group. +""" + +import pytest + +from modalities.exceptions import ConfigError +from modalities.running_env.fsdp.device_mesh import DeviceMeshConfig + + +def _config(**overrides) -> DeviceMeshConfig: + kwargs = dict(data_parallel_shard_degree=-1, world_size=8) + return DeviceMeshConfig(**{**kwargs, **overrides}) + + +def test_expert_parallel_degree_defaults_to_one(): + assert _config().expert_parallel_degree == 1 + + +def test_expert_parallelism_does_not_consume_world_size(): + # dp_shard resolves to the full world size even though 4 ranks' worth of it carries the experts. + config = _config(data_parallel_shard_degree=-1, expert_parallel_degree=4) + assert config.data_parallel_shard_degree == 8 + assert config.expert_parallel_degree == 4 + + +@pytest.mark.parametrize("expert_parallel_degree", [1, 2, 4, 8]) +def test_expert_parallel_degree_dividing_dp_shard_is_accepted(expert_parallel_degree: int): + _config(data_parallel_shard_degree=8, expert_parallel_degree=expert_parallel_degree) + + +def test_expert_parallel_degree_exceeding_dp_shard_is_rejected(): + with pytest.raises(ConfigError, match="must not exceed data_parallel_shard_degree"): + _config(data_parallel_shard_degree=2, world_size=2, expert_parallel_degree=4) + + +def test_expert_parallel_degree_not_dividing_dp_shard_is_rejected(): + with pytest.raises(ConfigError, match="must be divisible by"): + _config(data_parallel_shard_degree=8, expert_parallel_degree=3) + + +@pytest.mark.parametrize( + "degree_name", ["tensor_parallel_degree", "pipeline_parallel_degree", "context_parallel_degree"] +) +def test_expert_parallelism_rejects_unsupported_combinations(degree_name: str): + with pytest.raises(ConfigError, match="not yet supported together with"): + _config(data_parallel_shard_degree=4, expert_parallel_degree=2, **{degree_name: 2}) diff --git a/tests/fsdp2_parallelization/nemotron_ep_fsdp2_config.yaml b/tests/fsdp2_parallelization/nemotron_ep_fsdp2_config.yaml new file mode 100644 index 000000000..8f1f654ca --- /dev/null +++ b/tests/fsdp2_parallelization/nemotron_ep_fsdp2_config.yaml @@ -0,0 +1,182 @@ +# 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 + expert parallelism tests. +# Exercises all four layer types (Mamba-2, MoE, attention, dense MLP) at a size that fits +# comfortably on two GPUs. +# +# Identical to nemotron_fsdp2_config.yaml except for `expert_parallel_degree` and the added +# `expert_parallelized_model` component, so the two configs form an A/B pair. + +device_mesh: + component_key: device_mesh + variant_key: default + config: + device_type: cuda + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + # Carved out of data_parallel_shard_degree, so the 8 routed experts are split over the ranks + # and nothing is left for FSDP to shard them over. + expert_parallel_degree: ${cuda_env:WORLD_SIZE} + 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] + +expert_parallelized_model: + component_key: model + variant_key: expert_parallelized + config: + model: + instance_key: model_raw + pass_type: BY_REFERENCE + device_mesh: + instance_key: device_mesh + pass_type: BY_REFERENCE + +activation_checkpointed_model: + component_key: model + variant_key: activation_checkpointed + config: + model: + instance_key: expert_parallelized_model + 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_expert_parallelism_fsdp2.py b/tests/fsdp2_parallelization/test_expert_parallelism_fsdp2.py new file mode 100644 index 000000000..6ba7bc4e3 --- /dev/null +++ b/tests/fsdp2_parallelization/test_expert_parallelism_fsdp2.py @@ -0,0 +1,247 @@ +"""Distributed tests for expert parallelism. + +Three things are checked, each of which only manifests with a real process group: + +* the routed expert weights end up as DTensors on the ``(dp_shard_mod_ep, ep)`` mesh, holding only + this rank's share of the experts, while the router and shared experts stay data-parallel, +* a full forward/backward/optimizer step runs, including gradient clipping, which has to combine + gradient norms across two different device meshes, +* an expert-parallel MoE layer computes exactly what a replicated one computes -- the dispatch and + combine all-to-all pair must be an identity on the routing result. +""" + +import os +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from pydantic import BaseModel +from torch.distributed.tensor import DTensor + +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 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.parallelism.expert_parallelism import ExpertParallelGroupedExperts, shard_experts_over_ep_mesh +from modalities.running_env.fsdp.device_mesh import ParallelismDegrees, get_device_mesh +from modalities.training.gradient_clipping.fsdp_gradient_clipper import FSDP2GradientClipper, GradientClippingMode +from tests.end2end_tests.custom_components import MultiProcessingCudaEnv + +CONFIG_PATH = Path(os.path.dirname(__file__)) / "nemotron_ep_fsdp2_config.yaml" +WORLD_SIZE = 2 +VOCAB_SIZE = 512 +SEQ_LEN = 32 +NUM_EXPERTS = 8 # must match the config + + +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 + device_mesh = components.device_mesh + + # The mesh must expose the EP dimensions and still resolve `dp_shard`, which is now a + # flattened alias rather than a named dimension. + assert device_mesh[ParallelismDegrees.EP.value].size() == WORLD_SIZE + assert device_mesh[ParallelismDegrees.DP_SHARD.value].size() == WORLD_SIZE + assert device_mesh[ParallelismDegrees.DP_SHARD_MOD_EP.value].size() == 1 + + 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: global shape unchanged, but the local shard holds only this rank's experts + # and lives on the two-dimensional expert mesh. + expert_weight = find("experts.experts.w1") + assert isinstance(expert_weight, DTensor) + assert expert_weight.shape[0] == NUM_EXPERTS, expert_weight.shape + assert expert_weight.device_mesh.mesh_dim_names == ( + ParallelismDegrees.DP_SHARD_MOD_EP.value, + ParallelismDegrees.EP.value, + ) + local_expert_weight = expert_weight.to_local() + assert local_expert_weight.shape[0] == NUM_EXPERTS // WORLD_SIZE, local_expert_weight.shape + + # The router is not expert-parallel: it stays on the plain data-parallel mesh. + gate = find("moe.router.gate.weight") + assert gate.device_mesh.mesh_dim_names == (ParallelismDegrees.DP_SHARD.value,) + + # Everything must be materialized off the meta device and finite. + for name, param in parameters.items(): + local = param.to_local() if isinstance(param, DTensor) 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" + + +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 + # Gradient clipping has to reduce norms over both the data-parallel and the expert mesh. + clipper = FSDP2GradientClipper( + model_parts=model, max_norm=1.0, norm_type=GradientClippingMode.P2_NORM, device_mesh=components.device_mesh + ) + + # Different data per rank, so the dispatch all-to-all actually moves tokens. + 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}) + batch = InferenceResultBatch(targets={"target_ids": targets}, predictions=predictions) + loss = loss_fn(batch) + assert torch.isfinite(loss), loss + loss.backward() + grad_norm = clipper.clip_gradients() + assert torch.isfinite(grad_norm), grad_norm + assert grad_norm > 0, grad_norm + optimizer.step() + optimizer.zero_grad(set_to_none=True) + losses.append(loss.item()) + + # Ranks must agree on the loss trend having moved; the values themselves differ by data. + assert losses[1] != losses[0] + + # Auxiliary-loss-free load balancing must still work under expert parallelism. The router + # counts tokens per *global* expert before dispatch, and the counts are reduced over every + # rank that holds a distinct data shard -- which, with expert parallelism carved out of + # dp_shard, is all of them. A bias stuck at zero would mean the reduction group is wrong. + moe_layers = [module for name, module in model.named_modules() if name.endswith(".moe")] + assert moe_layers, "no MoE layer found" + for moe in moe_layers: + expert_bias = moe.router.expert_bias + assert expert_bias.shape == (NUM_EXPERTS,), expert_bias.shape + assert not torch.all(expert_bias == 0), "expert bias was never updated" + # The counters are reset by the hook after each step. + assert torch.all(moe.router.tokens_per_expert == 0), moe.router.tokens_per_expert + # Every rank must arrive at the same bias, since the counts were all-reduced. + gathered = [torch.zeros_like(expert_bias) for _ in range(WORLD_SIZE)] + dist.all_gather(gathered, expert_bias) + torch.testing.assert_close(gathered[0], gathered[-1], rtol=0, atol=0) + + +def _equivalence_worker(process_id: int, tmp_path: str, rdvz_port: int): + """An expert-parallel MoE must match a replicated MoE bit-for-bit in the forward pass.""" + del tmp_path + with MultiProcessingCudaEnv( + process_group_backend=ProcessGroupBackendType.nccl, + global_rank=process_id, + local_rank=process_id, + world_size=WORLD_SIZE, + rdvz_port=rdvz_port, + ): + n_embd, ffn_hidden, top_k = 64, 128, 2 + device_mesh = get_device_mesh( + device_type="cuda", + data_parallel_replicate_degree=1, + data_parallel_shard_degree=WORLD_SIZE, + tensor_parallel_degree=1, + pipeline_parallel_degree=1, + context_parallel_degree=1, + enable_loss_parallel=False, + world_size=WORLD_SIZE, + expert_parallel_degree=WORLD_SIZE, + ) + ep_mesh = device_mesh[ParallelismDegrees.EP.value] + + def build() -> MoE: + torch.manual_seed(1234) # identical on every rank + return ( + MoE( + router=TopKRouter(n_embd=n_embd, num_experts=NUM_EXPERTS, top_k=top_k, route_scale=1.0), + experts=GroupedExperts( + n_embd=n_embd, + ffn_hidden=ffn_hidden, + num_experts=NUM_EXPERTS, + backend=ExpertsBackend.GROUPED_MM, + ), + shared_experts=None, + aux_loss_coeff=0.0, + ) + .cuda() + .to(torch.bfloat16) + ) + + replicated = build() + reference_w1 = replicated.experts.w1.detach().clone() + reference_w2 = replicated.experts.w2.detach().clone() + + expert_parallel = build() + shard_experts_over_ep_mesh(expert_parallel.experts, ep_mesh) + local_experts = NUM_EXPERTS // WORLD_SIZE + lo, hi = process_id * local_experts, (process_id + 1) * local_experts + with torch.no_grad(): + expert_parallel.experts.w1.to_local().copy_(reference_w1[lo:hi]) + expert_parallel.experts.w2.to_local().copy_(reference_w2[lo:hi]) + expert_parallel.experts = ExpertParallelGroupedExperts(experts=expert_parallel.experts, ep_mesh=ep_mesh) + + # Per-rank distinct input so routing differs and the all-to-all is non-trivial. + torch.manual_seed(999 + process_id) + x = torch.randn(2, 64, n_embd, device="cuda", dtype=torch.bfloat16) + + torch.testing.assert_close(expert_parallel(x), replicated(x), rtol=0, atol=0) + + # Gradients must match too, up to bf16 accumulation order. The replicated run only sees this + # rank's tokens, so its expert gradients are summed over ranks before comparing against the + # expert-parallel run, whose local experts saw every rank's tokens. + replicated(x).float().pow(2).mean().backward() + expert_parallel(x).float().pow(2).mean().backward() + reference_grad = replicated.experts.w1.grad.float() + dist.all_reduce(reference_grad) + local_grad = expert_parallel.experts.experts.w1.grad + local_grad = (local_grad.to_local() if isinstance(local_grad, DTensor) else local_grad).float() + torch.testing.assert_close(local_grad, reference_grad[lo:hi], rtol=1e-2, atol=1e-4) + + +@pytest.mark.skipif( + torch.cuda.device_count() < WORLD_SIZE, reason=f"expert parallelism test requires {WORLD_SIZE} GPUs" +) +@pytest.mark.parametrize( + "worker, rdvz_port", + [(_sharding_worker, 22431), (_training_step_worker, 22432), (_equivalence_worker, 22433)], + ids=["sharding", "training_step", "equivalence_to_replicated"], +) +def test_expert_parallelism(worker, rdvz_port, tmp_path): + mp.spawn(worker, args=(str(tmp_path), rdvz_port), nprocs=WORLD_SIZE, join=True) diff --git a/tests/models/nemotron/test_expert_parallelism.py b/tests/models/nemotron/test_expert_parallelism.py new file mode 100644 index 000000000..deffb9f2e --- /dev/null +++ b/tests/models/nemotron/test_expert_parallelism.py @@ -0,0 +1,106 @@ +"""Unit tests for the expert-parallel dispatch permutation. + +The permutation that regroups received tokens from sender-major to expert-major order is the one +piece of expert parallelism that is pure index arithmetic, so it is tested here without a process +group. The distributed behaviour is covered by +``tests/fsdp2_parallelization/test_expert_parallelism_fsdp2.py``. +""" + +import pytest +import torch + +from modalities.models.parallelism.expert_parallelism import _build_permute_indices + + +def _reference_permute_indices(counts: torch.Tensor) -> torch.Tensor: + """Straightforward loop reference: concatenate the (expert, sender) ranges in expert-major order. + + Args: + counts (torch.Tensor): Received token counts of shape ``(ep_degree, num_local_experts)``. + + Returns: + torch.Tensor: The expected gather indices. + """ + ep_degree, num_local_experts = counts.shape + flat = counts.reshape(-1) + starts = (flat.cumsum(0) - flat).view(ep_degree, num_local_experts) + ranges = [ + torch.arange(int(starts[sender, expert]), int(starts[sender, expert]) + int(counts[sender, expert])) + for expert in range(num_local_experts) + for sender in range(ep_degree) + ] + return torch.cat(ranges) if ranges else torch.zeros(0, dtype=torch.long) + + +def _build(counts: torch.Tensor) -> torch.Tensor: + ep_degree, num_local_experts = counts.shape + return _build_permute_indices( + tokens_per_expert_group=counts.reshape(-1), + ep_degree=ep_degree, + num_local_experts=num_local_experts, + total_tokens=int(counts.sum()), + ) + + +@pytest.mark.parametrize( + "counts", + [ + pytest.param(torch.tensor([[3, 1], [2, 4]]), id="balanced_2x2"), + pytest.param(torch.tensor([[0, 5], [7, 0]]), id="some_experts_empty"), + pytest.param(torch.tensor([[0, 0], [0, 0]]), id="all_empty"), + pytest.param(torch.tensor([[1]]), id="degenerate_single_block"), + pytest.param(torch.tensor([[2, 0, 3, 1], [0, 0, 0, 0], [1, 1, 1, 1]]), id="one_sender_silent"), + ], +) +def test_permute_indices_match_loop_reference(counts: torch.Tensor): + torch.testing.assert_close(_build(counts), _reference_permute_indices(counts)) + + +def test_permute_indices_are_a_permutation(): + torch.manual_seed(0) + counts = torch.randint(0, 9, (4, 6)) + indices = _build(counts) + total = int(counts.sum()) + assert indices.shape == (total,) + # A valid gather permutation visits every row of the receive buffer exactly once. + torch.testing.assert_close(indices.sort().values, torch.arange(total)) + + +def test_permute_indices_group_tokens_by_expert(): + """The permuted order must place all tokens of one expert in one contiguous block.""" + counts = torch.tensor([[2, 1], [3, 4]]) # 2 senders, 2 local experts + indices = _build(counts) + # Sender-major receive buffer labelled by the expert each row belongs to. + expert_of_row = torch.tensor([0, 0, 1, 0, 0, 0, 1, 1, 1, 1]) + permuted = expert_of_row[indices] + torch.testing.assert_close(permuted, torch.tensor([0, 0, 0, 0, 0, 1, 1, 1, 1, 1])) + + +def test_permute_indices_do_not_read_values_back_to_host(monkeypatch): + """The construction must not read any tensor value back to the host. + + A device-to-host copy per (sender, local expert) pair is what makes the obvious implementation + expensive: it costs thousands of synchronizations per step for a deep MoE model. Only + ``total_tokens`` may come from the host, and the caller passes it in because it already has it + from the all-to-all split lists. + """ + counts = torch.randint(0, 9, (4, 8)) + total = int(counts.sum()) + + observed: list[str] = [] + for name in ("item", "tolist", "numpy"): + original = getattr(torch.Tensor, name) + + def spy(self, *args, _name=name, _original=original, **kwargs): + observed.append(_name) + return _original(self, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, name, spy) + + _build_permute_indices( + tokens_per_expert_group=counts.reshape(-1), + ep_degree=4, + num_local_experts=8, + total_tokens=total, + ) + assert observed == [], f"permutation construction read values back to the host via {set(observed)}"