Description
torchao/float8/float8_linear.py backward crashes with RuntimeError: expected scalar type BFloat16 but found Float when backward scaling is ScalingType.DISABLED and the model runs under torch.autocast(dtype=bfloat16).
Line 128-129 — DISABLED branch uses saved weight_hp_t directly:
elif c.cast_config_weight_for_grad_input.scaling_type is ScalingType.DISABLED:
weight_t_maybe_fp8_dim0 = weight_hp_t # float32 (master weight)Line 143-145 — torch.mm with mismatched dtypes:
grad_input = torch.mm(
grad_output_reshaped_maybe_fp8_dim0, # bf16 (from autograd under autocast)
weight_t_maybe_fp8_dim0.t(), # float32 ← mismatch
)Under autocast, Float8Linear.forward receives input as bf16 and self.weight.t() as float32 (master weight). Both are saved via ctx.save_for_backward(input_hp, weight_hp_t) — a (bf16, float32) pair. When backward scaling is active, hp_tensor_to_float8_dynamic casts both to float8, so the mismatch is masked. When DISABLED, the raw saved tensors go directly to torch.mm, which does not perform automatic dtype promotion.
Proposed fix
# Line 128-129: align dtype with grad_output
elif c.cast_config_weight_for_grad_input.scaling_type is ScalingType.DISABLED:
weight_t_maybe_fp8_dim0 = weight_hp_t.to(
grad_output_reshaped_maybe_fp8_dim0.dtype
).to(dtype) is a no-op when dtypes already match (no autocast → both float32), so this has zero cost outside the autocast + DISABLED scenario. The grad_weight path is unaffected — both its operands (grad_output and input_hp) are bf16 under autocast.
Steps to reproduce
import torch
from torchao.float8 import Float8LinearConfig, CastConfig, ScalingType
from torchao.float8.float8_linear import Float8Linear
# Explicitly DISABLE all backward scaling
disabled = CastConfig(scaling_type=ScalingType.DISABLED)
config = Float8LinearConfig(
cast_config_grad_output=disabled,
cast_config_weight_for_grad_input=disabled,
cast_config_input_for_grad_weight=disabled,
cast_config_grad_output_for_grad_weight=disabled,
)
layer = Float8Linear(64, 64, config=config).cuda()
x = torch.randn(8, 64, device="cuda", requires_grad=True)
with torch.autocast("cuda", dtype=torch.bfloat16):
y = layer(x)
y.sum().backward()
# → RuntimeError: expected scalar type BFloat16 but found Float
Without autocast (both operands float32), the same config works fine — confirming this is specifically a mixed-precision dtype alignment issue.Environment
- torch version: 2.12.0.dev (any recent version)
- torchao version: 0.18.0 (any version with `torchao.float8.float8_linear`)
- Hardware: reproduces on CUDA (no special hardware needed); affects any backend using `ScalingType.DISABLED` backward with autocast
- Files: `torchao/float8/float8_linear.py` line 128-129 (1 location, ~3 lines fix)
Description
torchao/float8/float8_linear.pybackward crashes withRuntimeError: expected scalar type BFloat16 but found Floatwhen backward scaling isScalingType.DISABLEDand the model runs undertorch.autocast(dtype=bfloat16).Line 128-129 — DISABLED branch uses saved
weight_hp_tdirectly:Line 143-145 —
torch.mmwith mismatched dtypes:Under autocast,
Float8Linear.forwardreceivesinputas bf16 andself.weight.t()as float32 (master weight). Both are saved viactx.save_for_backward(input_hp, weight_hp_t)— a (bf16, float32) pair. When backward scaling is active,hp_tensor_to_float8_dynamiccasts both to float8, so the mismatch is masked. When DISABLED, the raw saved tensors go directly totorch.mm, which does not perform automatic dtype promotion.Proposed fix
.to(dtype)is a no-op when dtypes already match (no autocast → both float32), so this has zero cost outside the autocast + DISABLED scenario. Thegrad_weightpath is unaffected — both its operands (grad_outputandinput_hp) are bf16 under autocast.Steps to reproduce
Environment