Fix: Handle LowLevelZeroPlugin with use_fp8=True

[Truong5724]

📌 Checklist before creating the PR

• I have created an issue for this PR for traceability
• The title follows the standard format: [doc/gemini/tensor/...]: A concise description
• I have added relevant tags if possible for better categorization
• I have installed pre-commit: pip install pre-commit && pre-commit install

🚨 Issue number

Fixed #6387

📝 What does this PR do?

Problem

• When running training with FP8 enabled under TorchDynamo (see main.py in the linked issue), execution fails during the forward/backward passes with a runtime error:

RuntimeError: shape '[32, 512]' is invalid for input of size 512.

• The issue above comes from incompatibility between FakeTensor (used by TorchDynamo) and torch._scaled_mm().

• FakeTensor only tracks metadata (shape, dtype, device) and requires all operations to be traceable without real data. However, _scaled_mm() is a low-level kernel that depends on real tensor data and does not support FakeTensor.

• As a result, TorchDynamo cannot correctly trace this operation, leading to runtime errors such as shape mismatch.

---

Solution

To ensure compatibility with TorchDynamo (FakeTensor), this PR removes the dependency on torch._scaled_mm() and replaces it with explicit dequantization + standard matrix multiplication (edit in fp8.py).

1. Forward pass

Original implementation:

out = torch._scaled_mm(
    x_fp8,
    ctx.w_fp8_t,
    bias=bias,
    out_dtype=ctx.out_dtype,
    scale_a=inv_scale_x,
    scale_b=inv_scale_w,
    use_fast_accum=True,
)[0]

Replaced with:

x_deq = x_fp8.to(ctx.out_dtype) * inv_scale_x
w_t_deq = ctx.w_fp8_t.to(ctx.out_dtype) * inv_scale_w

out = x_deq @ w_t_deq
if bias is not None:
    out = out + bias.to(ctx.out_dtype)

---

2. Backward pass

Original implementation:

x_grad = torch._scaled_mm(
    out_grad_fp8,
    ctx.w_fp8_t.contiguous().t(),
    out_dtype=ctx.out_dtype,
    scale_a=out_grad_scale,
    scale_b=ctx.inv_scale_w,
    use_fast_accum=True,
)[0]

w_grad = torch._scaled_mm(
    out_grad_fp8.t().contiguous(),
    ctx.x_fp8.t().contiguous().t(),
    out_dtype=ctx.out_dtype,
    scale_a=out_grad_scale,
    scale_b=ctx.inv_scale_x,
    use_fast_accum=True,
)[0]

Replaced with:

out_grad_deq = (out_grad_fp8.to(ctx.out_dtype) * out_grad_scale).contiguous()
w_t_deq = (ctx.w_fp8_t.to(ctx.out_dtype) * ctx.inv_scale_w).contiguous()
x_deq = (ctx.x_fp8.to(ctx.out_dtype) * ctx.inv_scale_x).contiguous()

x_grad = out_grad_deq @ w_t_deq.t()
w_grad = out_grad_deq.t() @ x_deq

---

Implications

• Removes reliance on _scaled_mm(), which is not FakeTensor-safe.
• Ensures compatibility with TorchDynamo execution.
• Uses standard matrix multiplication for correctness and stability.

---

Verification

• Training runs successfully under TorchDynamo.
• No runtime errors related to FakeTensor or _scaled_mm().
• Numerical behavior remains consistent.

---

💥 Checklist before requesting a review

• I have linked my PR to an issue.
• My issue clearly describes the problem.
• I have performed a self-review of my code.
• I have added thorough tests.
• I have added docstrings for all implemented functions.

⭐️ Do you enjoy contributing to Colossal-AI?

• 🌝 Yes, I do.
• 🌚 No, I don't.

[JazyJiang]

Thanks for looking into the TorchDynamo issue with the FP8 linear layer. Unfortunately we can't take this approach: dequantizing to the output dtype and using a regular matmul removes the FP8 compute path entirely, so every FP8 user loses the speed and memory benefit of _scaled_mm, not only the LowLevelZeroPlugin case.

The root cause is that newer PyTorch versions return a Tensor instead of a tuple from torch._scaled_mm, and Inductor expects 2-D scale tensors. A fix that keeps the FP8 kernel path (handling both return types and reshaping the scales to (1, 1)) is being landed as part of our PyTorch upgrade work, and we verified it on H20 with torch 2.11.

Closing this in favor of that fix. Thanks again for the report!