[MX] Rowwise `W` cached for backwards

[jeromeku]

Describe the bug

In transformer_engine.pytorch.Linear, both rowwise and columnwise quantized W are saved for backwards.

However, only columnwise W is needed for backwards for mx mixed precision.

Notably, rowwise is discarded and only columnwise is saved in the newer transformer_engine.pytorch.ops.BasicLinear.

[coderabbitai]

📝 CodeRabbit Plan Mode

Generate an implementation plan and prompts that you can use with your favorite coding agent.

• Create Plan

Examples

• Example 1
• Example 2

---

🧪 Issue enrichment is currently in open beta.

You can configure auto-planning by selecting labels in the issue_enrichment configuration.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord!