diff --git a/tests/functional_tests/parallelism/L2_Parallelism_Qwen3_5MoE_EP2_Parity.sh b/tests/functional_tests/parallelism/L2_Parallelism_Qwen3_5MoE_EP2_Parity.sh new file mode 100644 index 0000000000..414a8a3231 --- /dev/null +++ b/tests/functional_tests/parallelism/L2_Parallelism_Qwen3_5MoE_EP2_Parity.sh @@ -0,0 +1,69 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +#!/bin/bash +# Expert-parallel parity for the Qwen3.5-MoE recipes, which ship ep_size 4-32. +# +# Both legs run on 2 ranks with dp_size=2 and differ only in ep_size, so the +# data sharding and FSDP wrapping are identical and the comparison isolates +# expert parallelism. That also allows a tighter bound than the PP test, whose +# single-rank baseline is not FSDP-wrapped. +# +# The proxy generates its own synthetic token sequences, so this test stages no +# tokenizer or dataset. +# +# Known gap: this catches EP changing the numbers, but not EP silently never +# being applied -- that would make both legs identical and pass. The PP test +# closes the equivalent hole by grepping for the static-metadata log line; +# `apply_ep` has no such line to grep. + +set -xeuo pipefail + +export PYTHONPATH=${PYTHONPATH:-}:$(pwd) +export CUDA_VISIBLE_DEVICES="0,1" + +RUN_DIR=$(mktemp -d) +cleanup() { rm -rf "$RUN_DIR"; } +trap cleanup EXIT + +COMMON_ARGS=( + --config tests/functional_tests/parallelism/qwen3_5_moe_proxy.yaml + --step_scheduler.max_steps 6 + --step_scheduler.global_batch_size 4 + --step_scheduler.local_batch_size 2 + --distributed.tp_size 1 + --distributed.cp_size 1 + --distributed.pp_size 1 +) + +# --- Reference: 2 ranks, data parallel only --- +TRANSFORMERS_OFFLINE=1 python -m torch.distributed.run --nproc_per_node=2 --nnodes=1 -m coverage run \ + examples/llm_finetune/finetune.py \ + "${COMMON_ARGS[@]}" \ + --checkpoint.checkpoint_dir "$RUN_DIR/dp2" \ + --distributed.ep_size 1 + +# --- Expert parallel: same 2 ranks, experts sharded --- +TRANSFORMERS_OFFLINE=1 python -m torch.distributed.run --nproc_per_node=2 --nnodes=1 -m coverage run \ + examples/llm_finetune/finetune.py \ + "${COMMON_ARGS[@]}" \ + --checkpoint.checkpoint_dir "$RUN_DIR/ep2" \ + --distributed.ep_size 2 + +python tests/functional_tests/parallelism/compare_parallel_parity.py \ + "$RUN_DIR/dp2/training.jsonl" \ + "$RUN_DIR/ep2/training.jsonl" \ + --axis ep \ + --loss-tol 0.02 \ + --grad-norm-rtol 0.05 diff --git a/tests/functional_tests/parallelism/L2_Parallelism_Qwen3_5MoE_PP2_Parity.sh b/tests/functional_tests/parallelism/L2_Parallelism_Qwen3_5MoE_PP2_Parity.sh new file mode 100644 index 0000000000..9a93a8f4a5 --- /dev/null +++ b/tests/functional_tests/parallelism/L2_Parallelism_Qwen3_5MoE_PP2_Parity.sh @@ -0,0 +1,86 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +#!/bin/bash +# Pipeline-parallel parity for the Qwen3.5-MoE recipes, which ship pp_size 2-8. +# +# Qwen3.5 alternates full and linear attention, so a stage boundary can land +# between two layer types that carry different state. The generic PP tests run +# a dense, uniform-layer model and do not cover that. +# +# The proxy generates its own synthetic token sequences, so this test stages no +# tokenizer or dataset. + +set -xeuo pipefail + +export PYTHONPATH=${PYTHONPATH:-}:$(pwd) +export CUDA_VISIBLE_DEVICES="0,1" + +RUN_DIR=$(mktemp -d) +LOG_FILE="$RUN_DIR/pp2.log" +cleanup() { rm -rf "$RUN_DIR"; } +trap cleanup EXIT + +COMMON_ARGS=( + --config tests/functional_tests/parallelism/qwen3_5_moe_proxy.yaml + --step_scheduler.max_steps 6 + --step_scheduler.global_batch_size 4 + --step_scheduler.local_batch_size 2 + --distributed.tp_size 1 + --distributed.cp_size 1 + --distributed.ep_size 1 +) + +# --- Baseline: single rank, no parallelism --- +TRANSFORMERS_OFFLINE=1 python -m torch.distributed.run --nproc_per_node=1 --nnodes=1 -m coverage run \ + examples/llm_finetune/finetune.py \ + "${COMMON_ARGS[@]}" \ + --checkpoint.checkpoint_dir "$RUN_DIR/baseline" \ + --distributed.pp_size 1 + +# --- Pipeline parallel: 2 ranks --- +TRANSFORMERS_OFFLINE=1 python -m torch.distributed.run --nproc_per_node=2 --nnodes=1 -m coverage run \ + examples/llm_finetune/finetune.py \ + "${COMMON_ARGS[@]}" \ + --checkpoint.checkpoint_dir "$RUN_DIR/pp2" \ + --distributed.pp_size 2 \ + 2>&1 | tee "$LOG_FILE" + +# Guard against the `_precompute_stage_shapes` bug from PR #2983. Assert the +# static path positively as well: if the precompute is skipped outright, the +# fallback log line disappears too and the negative grep alone would pass. +if grep -Eiq "dynamic .*metadata inference" "$LOG_FILE"; then + echo "ERROR: pipeline stages fell back to dynamic metadata inference instead of static metadata" + exit 1 +fi +if ! grep -q "Precomputed pipeline stage shapes" "$LOG_FILE"; then + echo "ERROR: pipeline stage shapes were never precomputed; static metadata did not run" + exit 1 +fi + +# Loss here is ~10.9 (random init over a 32k vocab) and both legs compute it in +# bf16, whose ~0.4% relative resolution at that magnitude is already ~0.04 +# absolute. 0.10 sits above that floor; a tighter absolute bound would be +# measuring bf16, not pipeline parallelism. +# +# Gradient norm is the sensitive half of this check and is compared relatively, +# so it is unaffected by the loss magnitude. The bound stays loose because the +# single-rank baseline runs unwrapped while the pp2 run goes through FSDP2's +# bf16 mixed-precision policy. +python tests/functional_tests/parallelism/compare_parallel_parity.py \ + "$RUN_DIR/baseline/training.jsonl" \ + "$RUN_DIR/pp2/training.jsonl" \ + --axis pp \ + --loss-tol 0.10 \ + --grad-norm-rtol 0.20 diff --git a/tests/functional_tests/parallelism/qwen3_5_moe_proxy.yaml b/tests/functional_tests/parallelism/qwen3_5_moe_proxy.yaml new file mode 100644 index 0000000000..cb93c9fd39 --- /dev/null +++ b/tests/functional_tests/parallelism/qwen3_5_moe_proxy.yaml @@ -0,0 +1,149 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Small randomly-initialized Qwen3.5-MoE for the 2-GPU runners, shrunk from +# `examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b_neat_packing.yaml` (pp2, ep4). +# +# That recipe is a VLM, but no Qwen3.5 processor is staged in TEST_DATA_DIR, so +# this uses the text-only `Qwen3_5MoeForCausalLM`. Expert routing, expert +# sharding and the pipeline split all live in the text part, so they are +# covered; the vision tower is not. +# +# The expert counts are not scaled down, since that is what EP shards. +# +# Backends: `attn` matches the recipe. `experts` is torch_mm because the +# recipe's `gmm` only works with a deepep dispatcher. `dispatcher` is torch +# because deepep falls back to plain experts at world_size 1, which would make +# the PP test's 1-rank baseline run different code than its 2-rank leg. + +recipe: TrainFinetuneRecipeForNextTokenPrediction + +step_scheduler: + global_batch_size: 4 + local_batch_size: 2 + ckpt_every_steps: 1000 + val_every_steps: 1000 + max_steps: 6 + +dist_env: + backend: nccl + timeout_minutes: 10 + +rng: + _target_: nemo_automodel.components.training.rng.StatefulRNG + seed: 42 + ranked: true + +model: + _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_config + config: + _target_: transformers.models.qwen3_5_moe.configuration_qwen3_5_moe.Qwen3_5MoeTextConfig + # Routes NeMoAutoModel to the NeMo class via MODEL_ARCH_MAPPING, not the HF one. + architectures: [Qwen3_5MoeForCausalLM] + vocab_size: 32000 + hidden_size: 256 + num_hidden_layers: 6 + num_attention_heads: 8 + num_key_value_heads: 2 + head_dim: 32 + intermediate_size: 256 + moe_intermediate_size: 128 + shared_expert_intermediate_size: 128 + num_experts: 8 + num_experts_per_tok: 2 + max_position_embeddings: 4096 + rms_norm_eps: 1.0e-6 + # pad_token_id is left unset on purpose: it makes the embedding carry a + # padding_idx, and under FSDP2 that makes weight init get skipped, so the + # model trains on uninitialized weights. + # All full_attention: linear-attention layers keep two params in fp32, and + # this class does not set up the separate FSDP group they need. + layer_types: [full_attention, full_attention, full_attention, full_attention, full_attention, full_attention] + dtype: bfloat16 + backend: + _target_: nemo_automodel.components.models.common.BackendConfig + attn: sdpa + linear: torch + rms_norm: torch_fp32 + rope_fusion: false + experts: torch_mm + dispatcher: torch + fake_balanced_gate: false + enable_hf_state_dict_adapter: false + # Recipe value. Gates the MoEFSDPSyncMixin sync deferral that the + # gradient-accumulation path under test runs through. + enable_fsdp_optimizations: true + +checkpoint: + enabled: false + checkpoint_dir: checkpoints/qwen3_5_moe_proxy/ + +distributed: + strategy: fsdp2 + dp_size: none + tp_size: 1 + cp_size: 1 + pp_size: 1 + ep_size: 1 + + sequence_parallel: false + activation_checkpointing: false + + pipeline: + pp_schedule: 1f1b + pp_microbatch_size: 1 + scale_grads_in_schedule: false + # Recipe values. + patch_inner_model: false + patch_causal_lm_model: false + +loss_fn: + _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy + +# Synthetic fixed-length token sequences, so the test depends on no staged +# tokenizer or dataset. What is under test is whether two topologies agree on +# identical inputs, not what those inputs say. `std_len: 0` makes every sample +# the same length, which is what lets the pipeline stages use static shapes. +dataset: + _target_: nemo_automodel.components.datasets.llm.mock.build_unpacked_dataset + num_sentences: 64 + mean_len: 127 + std_len: 0 + vocab_size: 32000 + max_sentence_len: 128 + seed: 42 + +dataloader: + _target_: torchdata.stateful_dataloader.StatefulDataLoader + collate_fn: nemo_automodel.components.datasets.utils.default_collater + shuffle: false + +validation_dataset: + _target_: nemo_automodel.components.datasets.llm.mock.build_unpacked_dataset + num_sentences: 8 + mean_len: 127 + std_len: 0 + vocab_size: 32000 + max_sentence_len: 128 + seed: 7 + +validation_dataloader: + _target_: torchdata.stateful_dataloader.StatefulDataLoader + collate_fn: nemo_automodel.components.datasets.utils.default_collater + +optimizer: + _target_: torch.optim.AdamW + lr: 1e-5 + weight_decay: 0.01 + betas: [0.9, 0.95] diff --git a/tests/functional_tests/parallelism/test_parallelism.py b/tests/functional_tests/parallelism/test_parallelism.py index 09a05db6b8..01acc138c2 100644 --- a/tests/functional_tests/parallelism/test_parallelism.py +++ b/tests/functional_tests/parallelism/test_parallelism.py @@ -34,6 +34,8 @@ PP_GRAD_ACCUM_PARITY_FILENAME = "L2_Parallelism_PP_Grad_Accum_Parity.sh" DEEPSEEK_V4_PP2_PARITY_FILENAME = "L2_Parallelism_DeepSeekV4_PP2_Parity.sh" DEEPSEEK_V4_EP2_PARITY_FILENAME = "L2_Parallelism_DeepSeekV4_EP2_Parity.sh" +QWEN3_5_MOE_PP2_PARITY_FILENAME = "L2_Parallelism_Qwen3_5MoE_PP2_Parity.sh" +QWEN3_5_MOE_EP2_PARITY_FILENAME = "L2_Parallelism_Qwen3_5MoE_EP2_Parity.sh" class TestParallelismParity: @@ -51,3 +53,9 @@ def test_deepseek_v4_pp2_parity(self): def test_deepseek_v4_ep2_parity(self): run_test_script(TEST_FOLDER, DEEPSEEK_V4_EP2_PARITY_FILENAME) + + def test_qwen3_5_moe_pp2_parity(self): + run_test_script(TEST_FOLDER, QWEN3_5_MOE_PP2_PARITY_FILENAME) + + def test_qwen3_5_moe_ep2_parity(self): + run_test_script(TEST_FOLDER, QWEN3_5_MOE_EP2_PARITY_FILENAME)