feat: Add context parallel support for Qwen3.5 MoE

[zpqiu]

What does this PR do ?

Add CP support for Qwen3.5 MoE.

Changelog

• Add specific line by line info of high level changes in this PR.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)
  Add context parallel support for Qwen3.5 #1498

Convergence Experiment

• Pink: CP1
• Orange: CP2

# Qwen3.5 MoE 35B LLM SFT — CP=2 experiment
# Usage:
#   torchrun --nproc-per-node=8 examples/llm_finetune/finetune.py \
#     --config examples/llm_finetune/qwen/qwen3_5_moe_35b_cp2.yaml

step_scheduler:
    global_batch_size: 128
    local_batch_size: 4
    ckpt_every_steps: 500
    val_every_steps: 10
    num_epochs: 1
    max_steps: 200

dist_env:
    backend: nccl
    timeout_minutes: 60

rng:
    _target_: nemo_automodel.components.training.rng.StatefulRNG
    seed: 1234
    ranked: true

model:
    _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
    pretrained_model_name_or_path: Qwen/Qwen3.5-35B-A3B-Base
    backend:
        _target_: nemo_automodel.components.models.common.BackendConfig
        attn: te
        linear: te
        rms_norm: torch_fp32
        experts: te
        dispatcher: no
        rope_fusion: false
        fake_balanced_gate: false
        enable_hf_state_dict_adapter: true

checkpoint:
    enabled: false

distributed:
    strategy: fsdp2
    tp_size: 1
    cp_size: 2
    pp_size: 1
    ep_size: 8

    sequence_parallel: false
    activation_checkpointing: true

loss_fn:
    _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy

dataset:
    _target_: nemo_automodel.components.datasets.llm.chat_dataset.ChatDataset
    path_or_dataset_id: allenai/tulu-3-sft-mixture
    split: train
    shuffle_seed: 1234
    seq_length: 2048
    padding: max_length
    truncation: true

packed_sequence:
    packed_sequence_size: 0

dataloader:
    _target_: torchdata.stateful_dataloader.StatefulDataLoader
    collate_fn: nemo_automodel.components.datasets.utils.default_collater
    shuffle: true

validation_dataset:
    _target_: nemo_automodel.components.datasets.llm.chat_dataset.ChatDataset
    path_or_dataset_id: allenai/tulu-3-sft-mixture
    split: "train[:1024]"
    shuffle_seed: 1234
    seq_length: 2048
    padding: max_length
    truncation: true

validation_dataloader:
    _target_: torchdata.stateful_dataloader.StatefulDataLoader
    collate_fn: nemo_automodel.components.datasets.utils.default_collater
    batch_size: 4

optimizer:
    _target_: torch.optim.Adam
    betas: [0.9, 0.999]
    eps: 1e-7
    lr: 1.0e-5
    weight_decay: 0
    foreach: false

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[zpqiu]

/ok to test 4807fb2

[copy-pr-bot]

/ok to test 4807fb2

@zpqiu, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[zpqiu]

/ok to test 8978884

[akoumpa]

/ok to test 8e70593

[akoumpa]

/ok to test b950cc9

[akoumpa]

Thanks a lot @zpqiu for the contribution, very much appreciated. Please feel free to LMK if there's anything I can help with.