Support power of 2 scaling factors in float8 training #1670

danielvegamyhre · 2025-02-05T20:02:34Z

Summary

Add support for power of 2 scaling factors in float8 training with dynamic scaling.

Behavior:

Default on in the rowwise scaling recipe in Float8LinearConfig returned from recipe_name_to_linear_config for rowwise scaling.
Default off for other cases.

Test Plan
Updated test cases to ensure power of 2 scaling does not impact numerics for axiswise dynamic scaling (eager and compiled)

pytorch-bot · 2025-02-05T20:02:38Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

ROCM Infra failures during checkout of PyTorch

✅ No Failures

As of commit ab93e18 with merge base 8afd10e ():
💚 Looks good so far! There are no failures yet. 💚

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

danielvegamyhre · 2025-02-05T21:46:52Z

torchao/float8/float8_scaling_utils.py

@@ -36,6 +36,7 @@ def hp_tensor_to_float8_dynamic(
    device_mesh=None,
    scaling_granularity: ScalingGranularity = ScalingGranularity.TENSORWISE,
    axiswise_dim: Optional[int] = None,
+    power_of_2_scale: bool = False,


Note for reviewer: this param list is getting pretty long, and 4 of the 9 params can be derived from the Float8LinearConfig. Any thoughts on refactoring to pass in the Float8LinearConfig directly?

sounds reasonable

Cool, I'll do that in a follow up so Less can begin scale testing after we merge this asap

vkuzo · 2025-02-05T22:20:34Z

torchao/float8/config.py

+    # reduce quantization error by avoiding rounding errors when multiplying/dividing
+    # by the scaling factor, as well as ensuring large values are quantized to the
+    # same value in the forward pass as the backward passes.
+    power_of_2_scale: bool = False


nit: maybe something like round_scales_to_power_of_2, to match naming of surrounding code?

vkuzo

looks great! commented on suggested naming, any chance we could also check that this does not regress performance in torchtitan?

danielvegamyhre · 2025-02-05T23:09:37Z

looks great! commented on suggested naming, any chance we could also check that this does not regress performance in torchtitan?

My shared devgpu has too much other usage currently to test with Llama3 8b, so I used a debug model which is much smaller (800M params).

For rowwise with power of 2 scales, memory is flat but there is an ~8% regression in TPS. I'm wondering if this is because a small model like this has higher performance variance or if there is a real issue, need to look into it further.

Llama3 model configs: n_layers=4, dim=4096, n_heads=16

Tested on 4 H100s.

Row wise without power of 2 scales:

[rank0]:2025-02-05 15:05:21,682 - root - INFO - step:  1  loss:  8.2016  memory: 10.28GiB(10.82%)  tps: 584  mfu: 0.31%
[rank0]:2025-02-05 15:05:21,682 - root - INFO - Synchronizing and adjusting timeout for all ProcessGroups to 0:01:40
[rank0]:2025-02-05 15:05:21,914 - root - INFO - step:  2  loss:  9.2027  memory: 11.67GiB(12.28%)  tps: 70,711  mfu: 38.00%
[rank0]:2025-02-05 15:05:22,146 - root - INFO - step:  3  loss:  8.4319  memory: 11.67GiB(12.28%)  tps: 70,914  mfu: 38.11%
[rank0]:2025-02-05 15:05:22,375 - root - INFO - step:  4  loss: 13.0116  memory: 11.67GiB(12.28%)  tps: 71,446  mfu: 38.40%
[rank0]:2025-02-05 15:05:22,604 - root - INFO - step:  5  loss: 10.0891  memory: 11.67GiB(12.28%)  tps: 71,662  mfu: 38.51%
[rank0]:2025-02-05 15:05:22,835 - root - INFO - step:  6  loss:  8.8140  memory: 11.67GiB(12.28%)  tps: 71,041  mfu: 38.18%
[rank0]:2025-02-05 15:05:23,068 - root - INFO - step:  7  loss:  7.9921  memory: 11.67GiB(12.28%)  tps: 70,531  mfu: 37.91%
[rank0]:2025-02-05 15:05:23,297 - root - INFO - step:  8  loss:  7.5519  memory: 11.67GiB(12.28%)  tps: 71,670  mfu: 38.52%
[rank0]:2025-02-05 15:05:23,525 - root - INFO - step:  9  loss:  7.4012  memory: 11.67GiB(12.28%)  tps: 71,808  mfu: 38.59%
[rank0]:2025-02-05 15:05:23,754 - root - INFO - step: 10  loss:  7.2013  memory: 11.67GiB(12.28%)  tps: 71,647  mfu: 38.51%

Row wise with power of 2 scales:

[rank0]:2025-02-05 15:02:52,539 - root - INFO - step:  1  loss:  8.2104  memory:  9.85GiB(10.37%)  tps: 1,981  mfu: 1.06%
[rank0]:2025-02-05 15:02:52,539 - root - INFO - Synchronizing and adjusting timeout for all ProcessGroups to 0:01:40
[rank0]:2025-02-05 15:02:52,792 - root - INFO - step:  2  loss:  9.2376  memory: 11.20GiB(11.79%)  tps: 64,845  mfu: 34.85%
[rank0]:2025-02-05 15:02:53,046 - root - INFO - step:  3  loss:  8.6284  memory: 11.20GiB(11.79%)  tps: 64,742  mfu: 34.80%
[rank0]:2025-02-05 15:02:53,297 - root - INFO - step:  4  loss: 11.2887  memory: 11.20GiB(11.79%)  tps: 65,266  mfu: 35.08%
[rank0]:2025-02-05 15:02:53,548 - root - INFO - step:  5  loss:  9.4400  memory: 11.20GiB(11.79%)  tps: 65,429  mfu: 35.16%
[rank0]:2025-02-05 15:02:53,800 - root - INFO - step:  6  loss:  8.5271  memory: 11.20GiB(11.79%)  tps: 65,117  mfu: 35.00%
[rank0]:2025-02-05 15:02:54,055 - root - INFO - step:  7  loss:  7.8088  memory: 11.20GiB(11.79%)  tps: 64,426  mfu: 34.63%
[rank0]:2025-02-05 15:02:54,305 - root - INFO - step:  8  loss:  7.4392  memory: 11.20GiB(11.79%)  tps: 65,452  mfu: 35.18%
[rank0]:2025-02-05 15:02:54,555 - root - INFO - step:  9  loss:  7.3227  memory: 11.20GiB(11.79%)  tps: 65,783  mfu: 35.35%
[rank0]:2025-02-05 15:02:54,805 - root - INFO - step: 10  loss:  7.0642  memory: 11.20GiB(11.79%)  tps: 65,620  mfu: 35.27%

vkuzo · 2025-02-05T23:19:25Z

sounds like that's worth a follow-up

two things to check I can think of:

does it reproduce on full size LLaMa 3 8B on 8 H100s?
does the regression go away if we use bit shifting instead of exp and log?

vkuzo · 2025-02-05T23:21:41Z

torchao/float8/float8_utils.py

-
+    if round_scales_to_power_of_2:
+        # rounds down to the nearest power of 2.
+        res = torch.exp2(torch.floor(torch.log2(res)))


this should be the same as setting the mantissa to all-zeroes (maybe with some special handling for inf/nan), and can be implemented with bit shifting. Do you want to try to see if that resolves the regression?

didn't test, but something like

for float32

res = res.view(torch.uint32_t) res = (res >> 23) << 23 res = res.view(torch.float)

uint32 doesn't support bitshift ops apparently so I had to use int32. unit tests pass though and TPS regression is gone. will the sign bit affect anything? I did some manual tests in the interpreter and rounding seemed to work as expecting.

[rank0]:2025-02-05 16:11:30,663 - root - INFO - step: 1 loss: 8.2105 memory: 9.69GiB(10.20%) tps: 610 mfu: 0.33% [rank0]:2025-02-05 16:11:30,663 - root - INFO - Synchronizing and adjusting timeout for all ProcessGroups to 0:01:40 [rank0]:2025-02-05 16:11:30,896 - root - INFO - step: 2 loss: 9.2258 memory: 11.02GiB(11.60%) tps: 70,207 mfu: 37.73% [rank0]:2025-02-05 16:11:31,129 - root - INFO - step: 3 loss: 8.5120 memory: 11.02GiB(11.60%) tps: 70,377 mfu: 37.82% [rank0]:2025-02-05 16:11:31,361 - root - INFO - step: 4 loss: 11.7253 memory: 11.02GiB(11.60%) tps: 70,885 mfu: 38.10% [rank0]:2025-02-05 16:11:31,591 - root - INFO - step: 5 loss: 9.3686 memory: 11.02GiB(11.60%) tps: 71,365 mfu: 38.35% [rank0]:2025-02-05 16:11:31,823 - root - INFO - step: 6 loss: 8.5610 memory: 11.02GiB(11.60%) tps: 70,634 mfu: 37.96% [rank0]:2025-02-05 16:11:32,059 - root - INFO - step: 7 loss: 7.7763 memory: 11.02GiB(11.60%) tps: 69,681 mfu: 37.45% [rank0]:2025-02-05 16:11:32,287 - root - INFO - step: 8 loss: 7.4649 memory: 11.02GiB(11.60%) tps: 71,963 mfu: 38.68% [rank0]:2025-02-05 16:11:32,517 - root - INFO - step: 9 loss: 7.2956 memory: 11.02GiB(11.60%) tps: 71,188 mfu: 38.26% [rank0]:2025-02-05 16:11:32,749 - root - INFO - step: 10 loss: 7.1085 memory: 11.02GiB(11.60%) tps: 70,748 mfu: 38.02%```

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Feb 5, 2025

danielvegamyhre changed the title ~~Support power of 2 scaling factors in float8 training~~ Support power of 2 scaling factors in float8 training via boolean param in Float8LinearConfig Feb 5, 2025

danielvegamyhre marked this pull request as draft February 5, 2025 20:03

danielvegamyhre added the topic: new feature Use this tag if this PR adds a new feature label Feb 5, 2025

support power of 2 scaling factors in float8 training

f2433b1

danielvegamyhre force-pushed the powerof2 branch from 4ff4aca to f2433b1 Compare February 5, 2025 21:21

danielvegamyhre closed this Feb 5, 2025

danielvegamyhre reopened this Feb 5, 2025

danielvegamyhre changed the title ~~Support power of 2 scaling factors in float8 training via boolean param in Float8LinearConfig~~ Support power of 2 scaling factors in float8 training Feb 5, 2025

danielvegamyhre marked this pull request as ready for review February 5, 2025 21:37

fix linter issues

a9fe17e

danielvegamyhre force-pushed the powerof2 branch from ecc23ae to a9fe17e Compare February 5, 2025 21:39

danielvegamyhre commented Feb 5, 2025

View reviewed changes

danielvegamyhre requested a review from vkuzo February 5, 2025 21:47

danielvegamyhre added 2 commits February 5, 2025 13:58

power of 2 scale in amax_to_scale

896bd8f

add docstring to tensor_to_scale

c70ad60

danielvegamyhre force-pushed the powerof2 branch from 067db27 to c70ad60 Compare February 5, 2025 22:02

vkuzo reviewed Feb 5, 2025

View reviewed changes

vkuzo approved these changes Feb 5, 2025

View reviewed changes

rename round_scales_to_power_of_2

34cc033

vkuzo reviewed Feb 5, 2025

View reviewed changes

use bitshifting for power of 2 rounding

ab93e18

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support power of 2 scaling factors in float8 training #1670

Support power of 2 scaling factors in float8 training #1670

danielvegamyhre commented Feb 5, 2025 •

edited

Loading

pytorch-bot bot commented Feb 5, 2025 •

edited

Loading

danielvegamyhre Feb 5, 2025

vkuzo Feb 5, 2025

danielvegamyhre Feb 5, 2025

vkuzo Feb 5, 2025

danielvegamyhre Feb 5, 2025

vkuzo left a comment

danielvegamyhre commented Feb 5, 2025 •

edited

Loading

vkuzo commented Feb 5, 2025

vkuzo Feb 5, 2025

vkuzo Feb 5, 2025

danielvegamyhre Feb 6, 2025

Support power of 2 scaling factors in float8 training #1670

Are you sure you want to change the base?

Support power of 2 scaling factors in float8 training #1670

Conversation

danielvegamyhre commented Feb 5, 2025 • edited Loading

pytorch-bot bot commented Feb 5, 2025 • edited Loading

🔗 Helpful Links

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

❗ 1 Active SEVs

✅ No Failures

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

vkuzo left a comment

Choose a reason for hiding this comment

danielvegamyhre commented Feb 5, 2025 • edited Loading

vkuzo commented Feb 5, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

for float32

Choose a reason for hiding this comment

danielvegamyhre commented Feb 5, 2025 •

edited

Loading

pytorch-bot bot commented Feb 5, 2025 •

edited

Loading

danielvegamyhre commented Feb 5, 2025 •

edited

Loading