Initial prototype of differentiable _scaled_grouped_mm function

[danielvegamyhre]

Summary

The _grouped_scaled_mm function in torchao will do:

• Dynamic, float8 rowwise quantization on inputs
• Use these float8 inputs with the grouped scaled mm kernel in pytorch core and return the result
• Do this in a differentiable way

Note this prototype only handles A=2D, B=3D.

Test plan

• Added unit tests verifying the correctness of the forward pass (outputs) and backward pass (gradients)
• Verified torch.compile works with no graph breaks

Example usage

from torchao.prototype.scaled_grouped_mm import _scaled_grouped_mm

...

out = _scaled_grouped_mm(
    x,             # 2D high precision input tensor
    params,        # 3D high precision weights
    offs=offs,     # 1D int32 group offsets
    out_dtype=out_dtype,
)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1969

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 300db8b with merge base 923242e ():

NEW FAILURE - The following job has failed:

• .github/workflows/float8nocompile_test.yaml (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[ngimel]

this is causing device-host synchronization, and we should avoid it, relying on upstream ops to create suitable inputs, or in worst case, on _scaled_grouped_mm implementation itself to throw an assert

[danielvegamyhre]

Good catch, removed this assertion for now and will rely on kernel side assert, to avoid device-host sync

[ngimel]

offs is also a tensor so better you you use save_for_backward

[danielvegamyhre]

Done

[ngimel]

It might be better to scale and transpose B in forward, and store only quantized version (to minimize memory)

[danielvegamyhre]

good idea, done. hopefully torch.compile can do some fusion here and read B once and write both outputs simultaneously (float transposed column major, float8 non-transposed column major).

[ngimel]

here also would be better to have a triton kernel computing scales that could read offs on the device, to avoid syncs

[danielvegamyhre]

Yeah this implementation has more room for perf optimization, my first goal was to get accurate numerics. As a follow up I can write a triton kernel to avoid this device-host sync and for loop.

[vkuzo]

it's suprising to see a config created here, why not just inline the logic you need without worrying about configs? IMO dealing with configs would be for when this API is about to be productionized.

[danielvegamyhre]

I agree, I used it here for this prototype because in the test code i need to use matmul_with_hp_or_float8_args (which requires a Float8 config) to compute the reference forward/backward - and trying to get the outputs/grads to match was already pretty tricky, so to start I wanted to just reference the same Float8LinearConfig.from_recipe_name(Float8LinearRecipeName.ROWWISE) everywhere, to minimize room for accidental differences in how a particular tensor is quantized in test code vs implementation, etc.

I've now updated the implementation to inline everything, then in the test code compare against the float8 rowwise recipe to verify correctness.

[vkuzo]

remove?

[vkuzo]

can we clarify if this is for A, B or both? A, right?

[vkuzo]

remove?

[vkuzo]

are these assertions redundant with what is there in torch._scaled_grouped_mm? Ideally we only assert in one place for each condition.

[danielvegamyhre]

Some are specific to this implementation (only support 2d A and 3D B for now, the primary use case), but some are redundant, yes.

I found the device-side assertions to be opaque sometimes, I often had to read the kernel code to figure out why it wasn't working, to see the exact condition that was failing if the error was a bit ambiguous. So my goal here was to make the requirements more transparent and easier to debug.

(I think since scaled grouped mm is a kernel executing on the GPU, if a check fails, then on the CPU side the error message we can't see the actual line of code with the condition that failed, it just points to the entrypoint torch._scaled_grouped_mm)

[ngimel]

ndim assertions are not device side, so they have correct stack trace. Device-side assertions are put in when it's impossible to check the same thing on the host without introducing host-device sync, so they happen for a reason and unfortunately can't be improved. To get correct stack trace run with CUDA_LAUNCH_BLOCKING=1

[vkuzo]

nit: I'd prefer letting the caller do this instead, and this function can just assert that the layout is what is needed for the kernel

[vkuzo]

this naming is confusing, if the original variable is B, then this looks like B_transposed. IMO it would be cleanest to do something like:

• input is B
• B transposed is B_t

or

• input is B_t
• B_t transposed is B

[danielvegamyhre]

Yeah I agree, I considered this as well. The problem is I'm trying to make the naming consistent with torch._scaled_grouped_mm, which calls the tensors A and B, but it checks that B must be transposed - although really what it's enforcing is column major format, so it's a bit confusing.

For now I changed the naming to option 2 above, which will make the python code here clearer, with the trade-off being it will no longer be consistent with the kernel naming. I think that's fine though, I doubt too many people will be diving into the kernel code.

[vkuzo]

does this work for 4d/5d/etc tensors? if not, maybe assert that rank is 3?