perf(magi): avoid TE padding for magi

examples/llm_finetune/qwen/qwen3_moe_2layer_magi_packed_cp8_32k.yaml

@@ -0,0 +1,132 @@
+# 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.
+
+# Single-node gate for long-context Qwen3-MoE Magi CP benchmarking. The model
+# keeps Qwen3-30B-A3B's production attention, router, expert, and vocabulary
+# dimensions, but uses two transformer layers so CP8 forward/backward/optimizer
+# tests finish quickly. DeepEP, TE experts, THD packing, and activation
+# checkpointing stay enabled. Compare the identical recipe with:
+#
+#   automodel <this-file> --nproc-per-node 8 --model.backend.attn=te
+#
+# Sweep 64K with --packed_sequence.packed_sequence_size=65536 and
+# --model.config.max_position_embeddings=65536 after the 32K gate passes.
+
+recipe: TrainFinetuneRecipeForNextTokenPrediction
+
+step_scheduler:
+  global_batch_size: 1
+  local_batch_size: 1
+  max_steps: 5
+  num_epochs: 1
+  ckpt_every_steps: 1000
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 20
+
+rng:
+  _target_: nemo_automodel.components.training.rng.StatefulRNG
+  seed: 1234
+  ranked: true
+
+model:
+  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_config
+  config:
+    _target_: transformers.models.qwen3_moe.configuration_qwen3_moe.Qwen3MoeConfig
+    architectures: [Qwen3MoeForCausalLM]
+    vocab_size: 151936
+    hidden_size: 2048
+    intermediate_size: 6144
+    moe_intermediate_size: 768
+    num_hidden_layers: 2
+    num_attention_heads: 32
+    num_key_value_heads: 4
+    head_dim: 128
+    max_position_embeddings: 32768
+    hidden_act: silu
+    rms_norm_eps: 1.0e-6
+    rope_theta: 1000000.0
+    num_experts: 128
+    num_experts_per_tok: 8
+    decoder_sparse_step: 1
+    mlp_only_layers: []
+    router_aux_loss_coef: 0.001
+    norm_topk_prob: true
+    initializer_range: 0.02
+    bos_token_id: 151643
+    eos_token_id: 151645
+    pad_token_id: 151643
+    use_cache: false
+  backend:
+    _target_: nemo_automodel.components.models.common.BackendConfig
+    attn: magi
+    linear: te
+    rms_norm: torch_fp32
+    experts: te
+    dispatcher: deepep
+    rope_fusion: false
+    fake_balanced_gate: false
+    enable_hf_state_dict_adapter: false
+
+checkpoint:
+  enabled: false
+  checkpoint_dir: /tmp/qwen3_magi_cp_smoke/
+  model_save_format: safetensors
+  save_consolidated: false
+
+distributed:
+  strategy: fsdp2
+  tp_size: 1
+  cp_size: 8
+  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.hellaswag.HellaSwag
+  path_or_dataset: rowan/hellaswag
+  split: train
+  num_samples_limit: 8192
+  pad_to_max_length: false
+  tokenizer:
+    pretrained_model_name_or_path: Qwen/Qwen3-30B-A3B
+
+packed_sequence:
+  packed_sequence_size: 32768
+
+dataloader:
+  _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  collate_fn: nemo_automodel.components.datasets.utils.packed_sequence_thd_collater
+  shuffle: false
+
+optimizer:
+  _target_: torch.optim.Adam
+  betas: [0.9, 0.999]
+  eps: 1.0e-7
+  lr: 1.0e-4
+  weight_decay: 0
+  foreach: false
+
+wandb:
+  enable: false
+
+ci:
+  recipe_owner: huiyingl
+  time: "00:20:00"
+  nodes: 1

nemo_automodel/recipes/llm/train_ft.py

@@ -719,12 +719,14 @@ def materialize_loader(config):
                     dp_world_size=self._get_dp_group_size(),
                     pp_enabled=self.pp_enabled,
                     supports_seq_lens=_supports_seq_lens(self.model_parts[0]),
-                    # Models that own their CP shard the packed row contiguously;
-                    # per-document CP padding is a TE blockdiag concern and would
-                    # make pack composition depend on the topology.
+                    # Models and backends that own their CP dispatch do not need
+                    # TE's per-document ``2 * cp_size`` packing alignment.
                     cp_size=(
                         1
-                        if getattr(self.model_parts[0], "_owns_cp_attention", False)
+                        if (
+                            getattr(self.model_parts[0], "_owns_cp_attention", False)
+                            or (self.magi.enabled and self.magi.custom)
+                        )
                         else self.cfg.get("distributed.cp_size", 1)
                     ),
                     attn_implementation=attn_implementation,