fix(update_weight): preserve grouped MoE expert axis during GLU rechunk

[LLMShark]

Background

The GLU rechunking logic in update_weight/common.py splits linear_fc1 tensors along dim=0 after TP all-gather to restore the gate/up partition order.

This is correct for:

• 2D dense or per-expert linear_fc1.weight tensors.
• 1D linear_fc1.bias tensors.

It is unsafe for grouped MoE tensors such as:

• Fused 3D weights with shape [num_experts, 2 * ffn, hidden].
• Fused 2D biases with shape [num_experts, 2 * ffn].

For grouped tensors, dim=0 is the expert axis. Applying chunk(2, dim=0) would split the experts instead of the gate/up partitions and corrupt the gathered weights.

Changes

Updated slime/backends/megatron_utils/update_weight/common.py:

• Added a shared _requires_glu_rechunk(name, tensor) predicate.
• GLU rechunking is applied only to:
  • 2D linear_fc1.weight tensors.
  • 1D linear_fc1.bias tensors.
• Grouped MoE tensors, including fused 3D weights and fused 2D biases, are preserved intact for downstream model-specific conversion.
• Both synchronous all_gather_param() and asynchronous all_gather_params_async() use the same predicate.

This prevents the synchronous and asynchronous weight-update paths from diverging while preserving the existing dense and per-expert GLU behavior.

Tests

Added tests/test_grouped_moe_glu_rechunk.py, covering:

• 2D linear_fc1.weight still requires GLU rechunking.
• 1D linear_fc1.bias still requires GLU rechunking.
• Fused 3D grouped-MoE linear_fc1.weight does not get rechunked.
• Fused 2D grouped-MoE linear_fc1.bias does not get rechunked.
• linear_fc2 is unaffected by GLU rechunking.
• Synchronous all-gather preserves the grouped-MoE expert axis.
• Asynchronous all-gather preserves the grouped-MoE expert axis.
• The existing 2D GLU partition ordering remains unchanged.

The test is included in the CPU CI matrix.

Validation

• git diff --check passed.
• Python syntax checks passed.
• CPU unit tests passed locally: 8 passed.
• GitHub PR Test passed, including test_grouped_moe_glu_rechunk.py.