Enabling TP on Qwen3-Next-80B-A3B will trigger error: `AssertionError: Tensor parallelism not supported for custom MoE models`

Describe the bug

Setting TP>=2, you will have the assertion error:

AssertionError: Tensor parallelism not supported for custom MoE models

Steps/Code to reproduce bug
docker image: nvcr.io/nvidia/nemo-automodel:25.11

yaml configuration:

step_scheduler:
  global_batch_size: 512
  local_batch_size: 8
  ckpt_every_steps: 500
  val_every_steps: 100  # will run every x number of gradient steps
  num_epochs: 1
  max_steps: 100
  # max_steps: 5

dist_env:
  backend: nccl
  timeout_minutes: 10

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

parallelizer:
  _target_: nemo_automodel.components.moe.parallelizer.parallelize_model
  activation_checkpointing: false

model:
  _target_: nemo_automodel.components.models.qwen3_next.model.Qwen3NextForCausalLM.from_pretrained
  #_target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
  pretrained_model_name_or_path: Qwen/Qwen3-Next-80B-A3B-Instruct
  backend:
    _target_: nemo_automodel.components.moe.utils.BackendConfig
    attn: te
    linear: te
    rms_norm: te
    enable_deepep: true
    fake_balanced_gate: false
    enable_hf_state_dict_adapter: true
    enable_fsdp_optimizations: true

checkpoint:
  enabled: false
  checkpoint_dir: ./checkpoints/Qwen3_Next_80B_A3B/
  model_save_format: safetensors # torch_save or safetensors
  save_consolidated: true # saves the model in a consolidated safetensors format. Requires model_save_format to be safetensors.

distributed:
  _target_: nemo_automodel.components.distributed.fsdp2.FSDP2Manager
  tp_size: 2 # enabling TP will trigger error: AssertionError: Tensor parallelism not supported for custom MoE models
  cp_size: 1
  pp_size: 2
  dp_replicate_size: 1
  ep_size: 8 #32
  sequence_parallel: false

autopipeline:
  _target_: nemo_automodel.components.distributed.pipelining.AutoPipeline
  pp_schedule: interleaved1f1b
  pp_microbatch_size: 2
  round_virtual_stages_to_pp_multiple: up
  scale_grads_in_schedule: false
  patch_inner_model: false
  patch_causal_lm_model: false


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

dataset:
  _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag
  path_or_dataset: rowan/hellaswag
  split: train

packed_sequence:
  # Set packed_sequence_size > 0 to run with packed sequences
  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.hellaswag.HellaSwag
  path_or_dataset: rowan/hellaswag
  split: validation
  num_samples_limit: 64

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

optimizer:
  _target_: torch.optim.Adam
  betas: [0.9, 0.95]
  eps: 1e-7
  lr: 1.0e-4
  weight_decay: 0.0
  # min_lr: 1.0e-5

# Uncomment and configure for W&B logging
# wandb:
#   project: <your_wandb_project>
#   entity: <your_wandb_entity>
#   name: <your_wandb_exp_name>
#   save_dir: <your_wandb_save_dir>

And it will generate the following errors on 4 nodes of 8*H100 GPUs:

```shell
l_Finetune/0 Keyword args       : {   'activation_checkpointing': False,
l_Finetune/0     'cp_axis_name': 'cp',
l_Finetune/0     'dp_axis_names': ('dp_shard_cp',),
l_Finetune/0     'ep_axis_name': 'ep',
l_Finetune/0     'ep_shard_axis_names': ('ep_shard',),
l_Finetune/0     'moe_mesh': DeviceMesh((pp=2, ep_shard=1, ep=8), device: 'cuda', stride: (8, 8, 1)),
l_Finetune/0     'pp_enabled': True,
l_Finetune/0     'tp_axis_name': 'tp',
l_Finetune/0     'world_mesh': DeviceMesh((pp=2, dp_replicate=1, dp_shard=8, cp=1, tp=2), device: 'cuda', stride: (16, 16, 2, 2, 1))}
l_Finetune/0 Exception          : Tensor parallelism not supported for custom MoE models
l_Finetune/0 
l_Finetune/0 [rank2]: Traceback (most recent call last):
l_Finetune/0 [rank2]:   File "/opt/Automodel/examples/llm_finetune/finetune.py", line 33, in <module>
l_Finetune/0 [rank2]:     main()
l_Finetune/0 [rank2]:   File "/opt/Automodel/examples/llm_finetune/finetune.py", line 28, in main
l_Finetune/0 [rank2]:     recipe.setup()
l_Finetune/0 [rank2]:   File "/opt/Automodel/nemo_automodel/recipes/llm/train_ft.py", line 976, in setup
l_Finetune/0 [rank2]:     model, model_state_dict_keys, self.optimizer, self.loss_fn, self.param_info = build_model_and_optimizer(
l_Finetune/0 [rank2]:                                                                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^
l_Finetune/0 [rank2]:   File "/opt/Automodel/nemo_automodel/recipes/llm/train_ft.py", line 251, in build_model_and_optimizer
l_Finetune/0 [rank2]:     autopipeline.build(model, loss_fn=loss_fn, parallelize_fn=parallelize_fn)
l_Finetune/0 [rank2]:   File "/opt/Automodel/nemo_automodel/components/distributed/pipelining/autopipeline.py", line 133, in build
l_Finetune/0 [rank2]:     pp_schedule_obj, model_parts, pp_has_first_stage, pp_has_last_stage, stages = pipeline_model(
l_Finetune/0 [rank2]:                                                                                   ^^^^^^^^^^^^^^^
l_Finetune/0 [rank2]:   File "/opt/Automodel/nemo_automodel/components/distributed/pipelining/functional.py", line 496, in pipeline_model
l_Finetune/0 [rank2]:     parallelize_fn(
l_Finetune/0 [rank2]:   File "/opt/Automodel/nemo_automodel/components/config/loader.py", line 378, in instantiate
l_Finetune/0 [rank2]:     return func(*args, **config_kwargs)
l_Finetune/0 [rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
l_Finetune/0 [rank2]:   File "/opt/Automodel/nemo_automodel/components/moe/parallelizer.py", line 256, in parallelize_model
l_Finetune/0 [rank2]:     assert tp_axis_name is None or world_mesh[tp_axis_name].size() == 1, (
l_Finetune/0 [rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
l_Finetune/0 [rank2]: AssertionError: Tensor parallelism not supported for custom MoE models
l_Finetune/0 Instantiation failed for `parallelize_model`
l_Finetune/0 Accepted signature : (model: torch.nn.modules.module.Module, world_mesh: torch.distributed.device_mesh.DeviceMesh, moe_mesh: torch.distributed.device_mesh.DeviceMesh | None, *, pp_enabled: bool, dp_axis_names: tuple[str, ...], cp_axis_name: str | None = None, tp_axis_name: str | None = None, ep_axis_name: str | None = None, ep_shard_axis_names: tuple[str, ...] | None = None, activation_checkpointing: bool = False, reshard_after_forward: bool = False, lm_head_precision: str | torch.dtype | None = None)

Start by reviewing the merged pull request #1087, then inspect nemo_automodel/components/moe/parallelizer.py around line 256 and the parallelize_model call in the reported traceback. Reproduce the configuration with distributed.tp_size set to 2. Done means Qwen3-Next-80B-A3B setup completes without the custom-MoE tensor-parallelism assertion.