Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

use nanobind ndarray to make use of bfloat16 weights possible (without coping the data) #5

Merged
merged 1 commit into from
Mar 19, 2025
Merged

use nanobind ndarray to make use of bfloat16 weights possible (without coping the data) #5

merged 1 commit into from
Mar 19, 2025

Conversation

DavidGinten
Copy link

coping the data)

@DavidGinten Graphite App
Copy link
Author

This stack of pull requests is managed by Graphite. Learn more about stacking.

@DavidGinten DavidGinten force-pushed the ginten/use_nanobind_ndarray_for_bfloat16_support branch 4 times, most recently from ac06f9b to a50d3b0 Compare February 28, 2025 10:52
@DavidGinten Graphite App
Copy link
Author

DavidGinten commented Feb 28, 2025 •
edited
Loading

First iteration, can compile and run with float16 and bfloat16. AFAIU we have now just one object which gets deleted at the end

@DavidGinten DavidGinten force-pushed the ginten/use_nanobind_ndarray_for_bfloat16_support branch from a50d3b0 to 8519a35 Compare February 28, 2025 16:47
@DavidGinten Graphite App
Copy link
Author

DavidGinten commented Feb 28, 2025 •
edited
Loading

How do we add a test for bfloat16 (currently we creat numby objects). Maybe we create the dlpack object directly. Or simply import torch, but this is then a further dependency. Will upstream accept? EDIT: has been clarified by now

@DavidGinten DavidGinten force-pushed the ginten/use_nanobind_ndarray_for_bfloat16_support branch from 8519a35 to 96aa072 Compare March 3, 2025 12:51
@DavidGinten DavidGinten marked this pull request as ready for review March 6, 2025 13:06
@DavidGinten DavidGinten requested a review from devtbi March 7, 2025 14:29
devtbi
devtbi reviewed Mar 10, 2025
View reviewed changes
chrsmcgrr
chrsmcgrr requested changes Mar 10, 2025
View reviewed changes
Copy link
Collaborator

@chrsmcgrr chrsmcgrr left a comment

Choose a reason for hiding this comment

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

I think we need to make sure ownership of the tensor is handled.

@DavidGinten DavidGinten force-pushed the ginten/use_nanobind_ndarray_for_bfloat16_support branch 3 times, most recently from c34f357 to f49e81e Compare March 11, 2025 09:12
@DavidGinten DavidGinten requested a review from chrsmcgrr March 11, 2025 09:18
chrsmcgrr
chrsmcgrr requested changes Mar 11, 2025
View reviewed changes
Copy link
Collaborator

@chrsmcgrr chrsmcgrr left a comment

Choose a reason for hiding this comment

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

I think we need to verify this change works.

chrsmcgrr
chrsmcgrr reviewed Mar 11, 2025
View reviewed changes
Copy link
Collaborator

@chrsmcgrr chrsmcgrr left a comment

Choose a reason for hiding this comment

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

Just looking into this a little more, the DLPack does not define semantics for transfer of ownership, so for now this ndarray function will have to assume the lifetime of the data is tied to the python interpreter. nanobinds should ensure that condition.

@DavidGinten DavidGinten force-pushed the ginten/use_nanobind_ndarray_for_bfloat16_support branch 4 times, most recently from 42c3cad to 893acad Compare March 17, 2025 16:03
@DavidGinten DavidGinten force-pushed the ginten/use_nanobind_ndarray_for_bfloat16_support branch from 893acad to e0ad43b Compare March 17, 2025 16:28
@DavidGinten DavidGinten requested a review from chrsmcgrr March 17, 2025 16:29
@chrsmcgrr chrsmcgrr requested a review from devtbi March 17, 2025 16:30
chrsmcgrr
chrsmcgrr approved these changes Mar 18, 2025
View reviewed changes
Copy link
Collaborator

@chrsmcgrr chrsmcgrr left a comment

Choose a reason for hiding this comment

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

Looks good to me. @devtbi can you give your opinion?

mlir/python/mlir/_mlir_libs/_mlir/ir.pyi Outdated
up to the caller to ensure that the buffer meets the characteristics
implied by the shape.

The backing buffer and any user objects will be retained for the lifetime
Copy link
Collaborator

Choose a reason for hiding this comment

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

I don't think this statement is true. NDArray is designed as a view of the memory. Ownership is not transferred AFAIK. Can you change the wording to reflect that?

And I believe this is also true for the original buffer constructor. AFAIK for numpy at least we can't transfer ownership of the memory.

Copy link
Author

@DavidGinten DavidGinten Mar 18, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

Regarding numpy: We transfer a memory view which 'Create a new memoryview object which references the given object.'

mlir/test/python/ir/array_attributes.py Outdated
self.dlpack_capsule = array.__dlpack__()

def __del__(self):
print("BACKING MEMORY DELETED")
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggested change
print("BACKING MEMORY DELETED")
print("DLPACK MEMORY DELETED")

we actually don't own the memory just the metadata about the tensor (dlpack).

@DavidGinten DavidGinten force-pushed the ginten/use_nanobind_ndarray_for_bfloat16_support branch from e0ad43b to 35d0d56 Compare March 18, 2025 13:35
@devtbi Graphite App
Copy link

devtbi commented Mar 19, 2025

can't find anything to complain about since my last review... looks good :)

@DavidGinten DavidGinten merged commit 8cf4acb into integrate-llvm-project-20250217 Mar 19, 2025
3 checks passed
chrsmcgrr pushed a commit that referenced this pull request Mar 19, 2025
@DavidGinten @chrsmcgrr
use nanobind ndarray to make use of bfloat16 weights possible (#5)
80d8dac 
(withoutcoping the data)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@chrsmcgrr chrsmcgrr chrsmcgrr approved these changes

@devtbi devtbi devtbi approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@DavidGinten @devtbi @chrsmcgrr