Thread the scaling type argument throughout fp8 #301

drisspg · 2024-07-03T00:05:23Z

Summary

This PR adds a ScalingGranularity Enum, and threads it though the stack to all the places we call 'tensor_to_amax" and tensor_to_scale.

Currently hardcodes TensroWise.Scaling in Float8Linear, Float8DynamicLinear, Float8InferenceLinear. Asserts that granularity is TensorWise for now.
Added this as a property of WeightWithDynamicFloat8CastTensor, since we need to know a prior how do do the scaling for fp8 comms.

Testing

============================================================================= test session starts =============================================================================
platform linux -- Python 3.12.4, pytest-7.4.0, pluggy-1.5.0
rootdir: /home/drisspg/meta/float8_experimental
plugins: hypothesis-6.104.1
collected 9 items                                                                                                                                                             

test/test_fsdp2/test_fsdp2_eager.py .........                                                                                                                           [100%]

============================================================================= 9 passed in 30.77s ==============================================================================
all tests successful

Stack from ghstack (oldest at bottom):

[ghstack-poisoned]

ghstack-source-id: a740bcf5ce2098160870ee8da085861e956c5254 Pull Request resolved: #301

[ghstack-poisoned]

ghstack-source-id: 705ded1417d32b52fec3bae871c5f0c2922a5d0e Pull Request resolved: #301

[ghstack-poisoned]

ghstack-source-id: 583ea3369732127d122e45e182e1d1bc7c45fcc0 Pull Request resolved: #301

[ghstack-poisoned]

ghstack-source-id: c34e19a3ce0453fd1abb2b05db6bbb60ce3c90b8 Pull Request resolved: #301

[ghstack-poisoned]

ghstack-source-id: 333db4234b522fa07eced1b63c3d998317955c74 Pull Request resolved: #301

drisspg · 2024-07-03T05:12:01Z

float8_experimental/float8_tensor.py

-        amax_buffer: Optional[torch.Tensor] = None,
-        mm_config: Optional[ScaledMMConfig] = None,
+        float8_dtype: torch.dtype,
+        amax_buffer: Optional[torch.Tensor],


I removed the defualt args since this is always called from inner func with defualt args

[ghstack-poisoned]

ghstack-source-id: aa6f0c03f3fffaee5277337518c520a5895719ba Pull Request resolved: #301

[ghstack-poisoned]

ghstack-source-id: b09361e159b17dafe7940b24b3482ed482bba811 Pull Request resolved: #301

.pre-commit-config.yaml

vkuzo · 2024-07-03T15:40:10Z

float8_experimental/float8_tensor.py

@@ -31,6 +28,20 @@
 )


+class ScalingStrategy(Enum):


thoughts about using Granularity, which is more specific than Strategy?

Yeah thats a better word, this needed some bikeshedding

vkuzo · 2024-07-03T15:41:20Z

float8_experimental/float8_tensor.py

-    return Float8Tensor(bits_fp8, x_scale, x.dtype, mm_config=mm_config)
+    return Float8Tensor(
+        bits_fp8,
+        x_scale,


just curious, since we decided to not add scaling_strategy to torch._scaled_mm, why do we need it here?

We could make this a property of Float8Tensor, e.g. infer from the existing scales... hmm
Actually I might like this more..
We need the enum still since we want modules to specify their granularity

[ghstack-poisoned]

ghstack-source-id: 6f9b9299f4429ede127c0ed639a652d8888e947a Pull Request resolved: #301

# Summary This PR adds a ScalingGranularity Enum, and threads it though the stack to all the places we call 'tensor_to_amax" and tensor_to_scale. - Currently hardcodes TensroWise.Scaling in Float8Linear, Float8DynamicLinear, Float8InferenceLinear. Asserts that granularity is TensorWise for now. - Added this as a property of WeightWithDynamicFloat8CastTensor, since we need to know a prior how do do the scaling for fp8 comms. ### Testing ``` Shell ============================================================================= test session starts ============================================================================= platform linux -- Python 3.12.4, pytest-7.4.0, pluggy-1.5.0 rootdir: /home/drisspg/meta/float8_experimental plugins: hypothesis-6.104.1 collected 9 items test/test_fsdp2/test_fsdp2_eager.py ......... [100%] ============================================================================= 9 passed in 30.77s ============================================================================== all tests successful ``` [ghstack-poisoned]

Update

9db7cdc

[ghstack-poisoned]

drisspg added a commit that referenced this pull request Jul 3, 2024

threading the needle

1487a6d

ghstack-source-id: a740bcf5ce2098160870ee8da085861e956c5254 Pull Request resolved: #301

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 Jul 3, 2024

drisspg added a commit that referenced this pull request Jul 3, 2024

Thread thought the scaling type argument to float8 constructors

fff5eab

ghstack-source-id: a740bcf5ce2098160870ee8da085861e956c5254 Pull Request resolved: #301

drisspg added a commit that referenced this pull request Jul 3, 2024

Thread through the scaling type argument to float8 constructors

30ef0be

ghstack-source-id: a740bcf5ce2098160870ee8da085861e956c5254 Pull Request resolved: #301

drisspg added a commit that referenced this pull request Jul 3, 2024

Thread through the scaling type argument to float8 constructors

1b0f6d9

ghstack-source-id: a740bcf5ce2098160870ee8da085861e956c5254 Pull Request resolved: #301

drisspg added a commit that referenced this pull request Jul 3, 2024

Thread through the scaling type argument to float8 constructors

41bc9a4

ghstack-source-id: a740bcf5ce2098160870ee8da085861e956c5254 Pull Request resolved: #301

drisspg added a commit that referenced this pull request Jul 3, 2024

Thread through the scaling type argument to float8 constructors

f66afac

ghstack-source-id: a740bcf5ce2098160870ee8da085861e956c5254 Pull Request resolved: #301

drisspg added a commit that referenced this pull request Jul 3, 2024

Thread through the scaling type argument to float8 constructors

52b7ab5

ghstack-source-id: a740bcf5ce2098160870ee8da085861e956c5254 Pull Request resolved: #301

drisspg added a commit that referenced this pull request Jul 3, 2024

Thread through the scaling type argument to float8 constructors

7ec86d5

ghstack-source-id: a740bcf5ce2098160870ee8da085861e956c5254 Pull Request resolved: #301

drisspg added a commit that referenced this pull request Jul 3, 2024

Thread through the scaling type argument to float8 constructors

108f2f3

ghstack-source-id: a740bcf5ce2098160870ee8da085861e956c5254 Pull Request resolved: #301

drisspg added a commit that referenced this pull request Jul 3, 2024

Thread through the scaling type argument to float8 constructors

a516ac5

ghstack-source-id: a740bcf5ce2098160870ee8da085861e956c5254 Pull Request resolved: #301

drisspg marked this pull request as draft July 3, 2024 00:10

Update

4fcd497

[ghstack-poisoned]

drisspg mentioned this pull request Jul 3, 2024

Add sanity checks to dtensor tests #302

Closed

drisspg added a commit that referenced this pull request Jul 3, 2024

Thread through the scaling type argument to float8 constructors

ba2b798

ghstack-source-id: 705ded1417d32b52fec3bae871c5f0c2922a5d0e Pull Request resolved: #301

Update

f7a67bb

[ghstack-poisoned]

drisspg added a commit that referenced this pull request Jul 3, 2024

Thread through the scaling type argument to float8 constructors

385202a

ghstack-source-id: 583ea3369732127d122e45e182e1d1bc7c45fcc0 Pull Request resolved: #301

Update

e9b5ab8

[ghstack-poisoned]

drisspg added a commit that referenced this pull request Jul 3, 2024

Thread through the scaling type argument to float8 constructors

f2e62c5

ghstack-source-id: c34e19a3ce0453fd1abb2b05db6bbb60ce3c90b8 Pull Request resolved: #301

drisspg changed the title ~~threading the needle~~ Thread through the scaling type argument throughout fp8 Jul 3, 2024

drisspg changed the title ~~Thread through the scaling type argument throughout fp8~~ Thread the scaling type argument throughout fp8 Jul 3, 2024

Update

a4c98c5

[ghstack-poisoned]

drisspg added a commit that referenced this pull request Jul 3, 2024

Thread through the scaling type argument to float8 constructors

6cef6a1

ghstack-source-id: 333db4234b522fa07eced1b63c3d998317955c74 Pull Request resolved: #301

drisspg commented Jul 3, 2024

View reviewed changes

Update

4e7184b

[ghstack-poisoned]

drisspg added a commit that referenced this pull request Jul 3, 2024

Thread through the scaling type argument to float8 constructors

288b752

ghstack-source-id: aa6f0c03f3fffaee5277337518c520a5895719ba Pull Request resolved: #301

Update

381018d

[ghstack-poisoned]

drisspg added a commit that referenced this pull request Jul 3, 2024

Thread through the scaling type argument to float8 constructors

73eadae

ghstack-source-id: b09361e159b17dafe7940b24b3482ed482bba811 Pull Request resolved: #301

drisspg marked this pull request as ready for review July 3, 2024 05:48

drisspg commented Jul 3, 2024

View reviewed changes

.pre-commit-config.yaml Show resolved Hide resolved

vkuzo reviewed Jul 3, 2024

View reviewed changes

Update

8404bf6

[ghstack-poisoned]

drisspg added a commit that referenced this pull request Jul 3, 2024

Thread through the scaling type argument to float8 constructors

97acc3c

ghstack-source-id: 6f9b9299f4429ede127c0ed639a652d8888e947a Pull Request resolved: #301

This was referenced Jul 3, 2024

Add utility for filtering out skipped tests in large cross-product groups #303

Open

Add rowwise scaling to Float8Inference module #305

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Thread the scaling type argument throughout fp8 #301

Thread the scaling type argument throughout fp8 #301

drisspg commented Jul 3, 2024 •

edited

Loading

drisspg Jul 3, 2024

vkuzo Jul 3, 2024

drisspg Jul 3, 2024

vkuzo Jul 3, 2024

drisspg Jul 3, 2024

Thread the scaling type argument throughout fp8 #301

Are you sure you want to change the base?

Thread the scaling type argument throughout fp8 #301

Conversation

drisspg commented Jul 3, 2024 • edited Loading

Summary

Testing

drisspg Jul 3, 2024

Choose a reason for hiding this comment

vkuzo Jul 3, 2024

Choose a reason for hiding this comment

drisspg Jul 3, 2024

Choose a reason for hiding this comment

vkuzo Jul 3, 2024

Choose a reason for hiding this comment

drisspg Jul 3, 2024

Choose a reason for hiding this comment

drisspg commented Jul 3, 2024 •

edited

Loading