GSoC 2024 Project 12: Prototype for JAX/Flax Models Support (or Keras 3, Mindspore) #23531
Replies: 4 comments 12 replies
-
|
@MonalSD, @awayzjj, @chux0519, @LucaTamSapienza, @AsakusaRinne, fyi |
Beta Was this translation helpful? Give feedback.
-
|
Hi Roman, I'm very interested in exploring this extended choice because I'm familiar with Tensorflow and XLA (As you can see, I'm one of the maintainers of Tensorflow.NET). Could you please talk about the rationale behind this idea more deeply? In the idea description, it's mentioned that "Jax/Flax provides much faster training than PyTorch/Tensorflow". However, I find this claim not entirely convincing because it's actually benefited from XLA compiler, which is unfair to be compared with native PyTorch code. (Please correct me if I'm mistaken here) 😊 I hope to be able to maintain this support for Jax/Flax in the long term if my prototype is accepted in the future, even after GSoC ends in September. So I would be more than happy if supporting Jax/Flax proves to be valuable for users and openvino team. would you mind if I send you an email to ask for some suggestions directly? The proposal submission deadline is approaching, and I believe I can complete the draft this week. :) Additionally, should I aim to complete the work within 175 hours? I'm not sure if more time is necessary to ensure a thorough completion. Thank you for your time and support. Best regards, |
Beta Was this translation helpful? Give feedback.
-
|
Hi Roman @rkazants Thank you a lot for your response! I intend to send you a proposal in this weekend. Before I do, could you please confirm whether the project is supposed to be completed within 175 hours? In my opinion, additional time might be required to ensure comprehensive completion. However, I will adhere to the 175-hour timeline in my proposal if that is the final decision. :) Additionally, would it be acceptable for me to submit proposals for two projects concurrently? Initially, my interest was in project 11, and I've finished the official proposal for it. However, this project with the prospect of integrating Jax/Flax support really attracts me so much. Although I am more drawn to this project, I am inclined to continue to submit my proposal for project 11, given the substantial preparation I've invested. Thank you for your consideration! Sincerely, |
Beta Was this translation helpful? Give feedback.
-
|
Hi @chux0519, You have the right exploration that we currently support JAX models by OV using JAX2TF. However, this tool has limitations and this is not direct and native support of JAX models. This approach requires intermediate model representation and serialization on a disk. We need a sort of support for conversion of JAX/Flax object from memory, when you pass it into Best regards, |
Beta Was this translation helpful? Give feedback.
-
|
@chux0519 Hi, Yongsheng, I'm glad to see you here. I think you're right about that it's possible to get HLO IR from Jax/PyTorch/Tensorflow. Please refer to pytorch-xla, tf2xla and open-xla. Initially xla is a part of tensorflow and it was extracted from it exactly because developers didn't want XLA (or HLO IR) could be used only in tensorflow. However, in OpenVINO, I don't think we should do similar things, to get HLO IR from models of all frameworks. Otherwise we're just repeating the same thing as open-xla does. For GSoC application, I view our relationship as not only competitive but also collaborative. Though I'll submit a proposal for it, I'd like to share what I know and find during this preparation. I welcome any opportunity for us to discuss and exchange ideas freely. :) Best regards, |
Beta Was this translation helpful? Give feedback.
-
|
Understood. Actually, I'm more familiar with Stable Diffusion, and I've already applied for Proposal 4. However, the work discussed in this thread seems more interesting and challenging. I'll continue to keep an eye on it, and I believe that even after successfully implementing the proof of concept (POC) this summer, there will still be plenty of things to do. Additionally, Rinne, as far as I know, individuals are allowed to submit two proposals simultaneously (assuming you have the time). I look forward to exchanging ideas and learning together with everyone. |
Beta Was this translation helpful? Give feedback.
-
|
@chux0519 Thank you for your reminder. I have previously reviewed the rules of GSoC, and my inquiry about acceptance is mainly to ask if mentors mind it, because only one proposal per student could be accepted finally. :) I'm more than excited if having opportunity to work along with you on Jax/Flax support in the future! I will actively share my works with you if my proposal for this project passes. |
Beta Was this translation helpful? Give feedback.
-
|
Hey @rkazants , @AsakusaRinne and @chux0519 , It's truly invigorating to witness such a wealth of innovative ideas being brought to the table for this project. As a newcomer to the world of open source, I find this level of enthusiasm incredibly motivating. I agree with Yongsheng's point regarding the necessity of preprocessing on the Python side before parsing and generating the corresponding IR. Developing an adapter for JAX-based models to standardize them into an abstract format seems like a good approach to start with. This would facilitate seamless processing and compatibility moving forward. Moreover, considering that locally saved models essentially adhere to a pytree structure, devising our reconstruction of this pytree during IR conversion would enable us to mirror the practices embraced by the wider community, thereby empowering new contributors with a familiar framework to engage with. Additionally, I propose that benchmarking our converted OpenVINO models against the original JAX models could serve as a valuable tool. Not only would it provide us with insights into the potential impact we can achieve, but it would also bolster the credibility of our proof of concept. Looking forward to delving deeper into these concepts with all of you. |
Beta Was this translation helpful? Give feedback.
-
|
Dear contributors (@MonalSD, @awayzjj, @chux0519, @LucaTamSapienza, @AsakusaRinne), If you want to apply for Best regards, |
Beta Was this translation helpful? Give feedback.
-
|
@Pranshu-S, please pay attention as well |
Beta Was this translation helpful? Give feedback.
-
|
@rkazants I have sent my application to your email yesterday and I will appreciate it if you would like to give me some suggestions before I submit it to the GSoC portal. :) |
Beta Was this translation helpful? Give feedback.
-
|
Hi @AsakusaRinne, I will take a look today later. Thanks! |
Beta Was this translation helpful? Give feedback.
-
|
Hi @rkazants I have mailed my application. It will be great to have your inputs. |
Beta Was this translation helpful? Give feedback.
-
|
@rkazants Hi, I have sent the application to your email. I would greatly appreciate any feedback . |
Beta Was this translation helpful? Give feedback.
-
Dear participants in GSoC'24,
Pay attention that project idea 12 was extended with a choice to implement Prototype for support of Flax/JAX models by OpenVINO.
Feel free to ask your questions under this thread.
Best regards,
Roman
Beta Was this translation helpful? Give feedback.
All reactions