Skip to content

Fix fp16/bf16 divmod and floor_divide precision (divide in float before rounding) - #4003

Closed
sashko-zakharchuk wants to merge 1 commit into
ml-explore:mainfrom
sashko-zakharchuk:fix-divmod-fp16-precision
Closed

Fix fp16/bf16 divmod and floor_divide precision (divide in float before rounding)#4003
sashko-zakharchuk wants to merge 1 commit into
ml-explore:mainfrom
sashko-zakharchuk:fix-divmod-fp16-precision

Conversation

@sashko-zakharchuk

Copy link
Copy Markdown
Contributor

On fp16/bf16, mx.divmod and mx.floor_divide divided x / y in the input precision
and narrowed the quotient to fp16/bf16 before trunc/floor. A true quotient just
below an integer rounds up when narrowed, so the result is one too high.

a = mx.array([86.5625], dtype=mx.float16)    # a / b = 2.99946
b = mx.array([28.859375], dtype=mx.float16)
mx.divmod(a, b)[0]        # before: 3    after: 2
mx.floor_divide(a, b)     # before: 3    after: 2

The fix computes the division in float32 for the half types before rounding: on CPU in
DivMod's trunc/fmod op, on Metal and CUDA in the FloorDivide functor that
divmod reuses, and in the floor_divide lowering in ops.cpp. float32, float64,
integer, and sign behavior are untouched (divmod still truncates). For non-negative
fp16/bf16 operands divmod and floor_divide now agree with each other, and with numpy
for same-dtype operands; for negatives divmod still truncates while floor_divide
floors.

This is also why test_divmod is flaky: it compares fp16 mx.divmod against numpy with
unseeded random inputs, and near an integer boundary they disagreed (about 0.04% of
fp16 pairs, more for bf16). With the division in float they match. The test gets a
deterministic boundary case for fp16 and bf16, plus a seeded bf16 batch checked against
a float32 reference (numpy has no bf16).

Part of #3994 (the fp16/bf16 precision part); the floored-vs-truncated sign defects are
a separate change.

Verification

  • CPU (Linux x86): fp16/bf16 divmod and floor_divide match a float-wide reference on
    0/200000 random pairs; 1e300 / 1e299 float64 stays 10 (no downcast); -7 / 3
    stays (-2, -1) (sign unchanged); full test_ops.py passes.
  • CUDA (RTX 5050, sm_120): fp16/bf16 divmod and floor_divide give 2 at the boundary,
    0/100000 divergence, sign unchanged.
  • Metal (M3 Ultra): fp16/bf16 divmod and floor_divide give 2 at the boundary,
    0/100000 divergence, sign unchanged; test_divmod passes.

@zcbenz zcbenz left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since this is going to add overhead I think we should follow what PyTorch does, which has the same behavior:

a = torch.tensor([86.5625], dtype=torch.float16)
b = torch.tensor([28.859375], dtype=torch.float16)

print(torch.div(a, b), torch.remainder(a, b))
tensor([3.], dtype=torch.float16) tensor([28.8438], dtype=torch.float16)

Generally when in doubt we would prefer speed over precision as long as there is a way for users to use high precision computation.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants