fix(model): normalize Qwen3-Omni integer masks for TE

[hbhflw2000]

Summary

• Normalize 2D integral Qwen3-Omni valid-token masks to the boolean
  masked-token format [B, 1, 1, S] expected by Transformer Engine.
• Preserve the existing behavior for boolean, packed-sequence, and
  unsupported mask formats.
• Add regression coverage for asymmetric QKV and gate/up tensor mappings,
  including expert-parallel export.

Problem

The processor and collator may emit 2D integer valid-token masks, while
the existing conversion only handled boolean masks. These integer masks
therefore bypassed normalization even though Transformer Engine expects
boolean masks with masked-token semantics.

The existing conversion tests also relied mainly on registry checks and
symmetric round trips, which could miss incorrect tensor ordering.

Scope

This follow-up is limited to:

• Qwen3-Omni attention-mask normalization.
• Regression tests for QKV, gate/up, and EP=2 expert mappings.

It does not change the production tensor-mapping implementation,
packed-sequence handling, conversion APIs, dependencies, or Megatron Core.

Validation

• 15 focused unit tests passed.
• Local 4-GPU pruned-model SFT completed one finite training step with
  zero skipped or NaN iterations.
• A 4-node / 32-GPU full-model SFT run with TP2, PP2, and EP8 completed
  one finite training step:
  • LM loss: 4.563817
  • Load-balancing loss: 2.480726
  • Skipped iterations: 0
  • NaN iterations: 0
  • All distributed ranks exited cleanly.

Related

Follow-up to #4988, addressing the post-merge review comments on integral
attention masks and semantic tensor-mapping coverage.

[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.

[yaoyu-33]

/ok to test 9801179

[kamran-nvidia]

/ok to test 0690c89