test(pp): add qwen3_5_moe PP and EP parity tests [3/4]

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

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

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]

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)