Grouped Linear issue

I am training a model using the Megatron-DeepSpeed framework. Previously, I was using DeepSpeed MoE (MoE layers implemented with DeepSpeed) and the training was numerically stable — the loss curve and gradient norms behaved as expected.

Recently, I replaced the DeepSpeed MoE implementation with Transformer Engine’s Grouped Linear / Grouped GEMM (for MoE experts). After switching, I observe two major issues:

1. Loss does not match the baseline (DeepSpeed MoE)

When running the same model, same configuration, same data, and same initialization:

The loss curve diverges from the baseline almost immediately.

Loss values cannot be aligned with the previous implementation even in the very first few steps.

Forward pass outputs appear numerically different from the original MoE implementation.

This suggests that the grouped GEMM path may be producing different numerical results compared to the standard GEMM/DeepSpeed MoE path.

2. Gradient norm becomes extremely large

After enabling TE Grouped Linear, I observe:

Very large grad norms, sometimes orders of magnitude larger than under the baseline implementation.

This seems to originate from the MoE layers using grouped GEMM.

When switching back to DeepSpeed MoE or standard Linear layers, the grad norm returns to normal.

This indicates a potential stability / scaling problem inside grouped GEMM.

Start at the Transformer Engine Grouped Linear and Grouped GEMM entry points used by the MoE experts, and compare their forward outputs and gradient norms with the DeepSpeed MoE baseline under identical configuration, data, and initialization. A useful resolution would include a minimal reproduction and identify the source of the numerical mismatch or gradient scaling difference.