[float8] improve eager numerics for dynamic scales and gets on par with torch.compile

[weifengpy]

Turning on/off float8 all-gather shows different numerics: #873

Root cause: hp_tensor_to_float8_dynamic have different numerics in eager vs compile

• compile does float32 upcasting at the begining of scale calculation, see tmp0 = tl.load(in_ptr0 + (x0), None).to(tl.float32) in triton kernels below
• eager does float32 upcasting at the end
• even if we do float32 upcasting at the begining for both compile and eager, triton has different numerics than eager. but float64 has bitwise on-par numerics

This PR does float64 upcasting to the begining of scale calculation

unit test: TORCH_LOGS="output_code" pytest -s test/float8/test_base.py -k test_dynamic_scale_parity

tmp0 = tl.load(in_ptr0 + (x0), None).to(tl.float32)
tmp1 = tl.load(in_ptr1 + (0))
tmp2 = tl.broadcast_to(tmp1, [XBLOCK])
tmp3 = tmp2.to(tl.float32)
tmp4 = tmp0 * tmp3
tmp5 = tmp4.to(tl.float32)
tmp6 = -448.0
tmp7 = triton_helpers.maximum(tmp5, tmp6)
tmp8 = 448.0
tmp9 = triton_helpers.minimum(tmp7, tmp8)
tmp10 = tmp9.to(tl.float32)
tmp11 = tmp10.to(tl.float8e4nv)
tl.store(out_ptr0 + (x0), tmp11, None)

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 3d0da20 with merge base b149edb ():
💚 Looks good so far! There are no failures yet. 💚

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

[weifengpy]

upcast amax in amax_to_scale instead of tensor_to_amax for 2 reasons

• we can still do bfloat16 all-reduce for amax
• safer to delayed scaling as it won't change dtype for amax_buffer

[vkuzo]

could you share why the upcasting happens?

[weifengpy]

I can look into inductor more on how it achieved fp64

[weifengpy]

torch.compile actually upcasts to float32 with tl.load(in_ptr0 + (x0), None).to(tl.float32). Upcasting to float64 further help because torch.compile and eager shows different numerics for 1.0 / float32 (but same numeric for float64)

The float32 numeric difference can be verified with

import torch
def upcast_reciprocal(inp: torch.Tensor):
    return inp.reciprocal()
inp = torch.full([], 0.00817871093750000000, device="cuda", dtype=torch.float32)
eager_scale = upcast_reciprocal(inp)
compile_scale = torch.compile(upcast_reciprocal)(inp)
fp64_ground_truth = inp.to(torch.float64).reciprocal()
assert torch.equal(eager_scale, compile_scale), f"{eager_scale=} vs {compile_scale=}, {fp64_ground_truth=}"

[weifengpy]

without the PR, the numerics looks like following
eager _scale=106.5000 vs compile _scale=106.1925...

after, eager is also 106.1925...

[awgu]

I am a bit surprised that we have to use fp64 in eager. Without the upcast to fp64, what is the dtype being used?

[weifengpy]

I am a bit surprised that we have to use fp64 in eager. Without the upcast to fp64, what is the dtype being used?

eager respect param dtype - if param is bf16, amax is bf16, scale compute is bf16. scale gets upcasted to fp32 at the very end

[weifengpy]

converting to draft as upcasting to torch.float32 in eager should be enough. digging into inductor as there might be nuances around how bfloat16 is handled

[weifengpy]

without upcasting, the eager numeric is like -157.00000000000000000000, compile is like -157.06507873535156250000

[weifengpy]

torch.compile upcast tensor ahead, see tmp0 = tl.load(in_ptr0 + (x0), None).to(tl.float32) in following output code

@triton.jit
def triton_(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 24576
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = tl.full([XBLOCK], True, tl.int1)
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (x0), None).to(tl.float32)

[vkuzo]

nit: move to test_compile.py since this is testing compile vs eager?

[vkuzo]

can we make this comment contain the context? something like

# Note: when the line below is compiled with `torch.compile`, `tensor` is automatically upcasted to `float32` to multiply with the scale
# In order to match numerics between eager and compile, we upcast manually here.

[vkuzo]

nit: can we describe in more detail why we are upcasting here

[vkuzo]

add comment to describe upcasting

[weifengpy]

improved comment

[vkuzo]

thanks for finding this, it would be great to have numerics match! At a high level, looks good. Could we add two more things,

1. verify that there is no performance regression with torch.compile
2. ensure tests pass in CI

[weifengpy]

thanks for finding this, it would be great to have numerics match! At a high level, looks good. Could we add two more things,

1. verify that there is no performance regression with torch.compile
2. ensure tests pass in CI

thanks for the feedback. will address them and publish for review

[weifengpy]

verified perf is neutral and make sure all CI passes

for perf, I ran benchmark_compiled_module(times=100, repeat=100) from the output code of pytest -s test/float8/test_compile.py -k test_dynamic_scale_numeric_parity. I tested on shape = (1024, 1024),

• without the change: 0.000035 secs P1612792187
• with the change: 0.000035 secs P1612791058

looking into the output code, the difrerence is mainly tmp3.to(tl.float64) or tmp3.to(tl.float32)

[weifengpy]

addressed feedback and is ready for review

[awgu]

dtype=torch.float64 only changes the accumulation dtype? if there is no noticeable cost to this, I wonder if we should be doing this in more places 🤔

[weifengpy]

good question! Actually I just updated the code to do _foreach_norm in original precision and do float64 upcasting before calculating scales. That ensures consistent implementation between prcompute and float8_utils.amax_to_scale

back to your question, I checked ForeachReduceOp.cu and it's dispatching to lpnorm_cleanup<scalar_t, NormType::LInf, out_t>. Not sure what's inside lpnorm_cleanup. But inf-norm is just max(abs) so not sure if they accumulate numerics
https://github.com/pytorch/pytorch/blob/a28b40fa74470058ca57d77652b9601bece2f4d5/aten/src/ATen/native/cuda/ForeachReduceOp.cu#L534-L535C19

[vkuzo]

this should be an object of type LinearMMConfig, I'm actually kind of surprised passing in Float8LinearConfig works :(

[weifengpy]

good catch! I switched to LinearMMConfig. Float8LinearConfig was working because I did not call matmul that requires access to self._linear_mm_config.

[vkuzo]

optional nit: upcast separately to make easier to read

amax_tensor = amax_tensor.to(torch.float64)
scale_tensor = ... / amax_tensor

[weifengpy]

updated as suggested