[RFC] MXNet 2.0 JVM Language development

[lanking520]

Since MXNet 2.0 development starts. I would like to initiate a discussion for the future development of JVM languages.

Proposal

1. Start cleaning on the existing APIs to adapt to 2.0
2. Start from ground up to rewrite the whole Scala/Java APIs
3. Start using DJL (djl.ai) as a frontend for MXnet JVM development
4. Using DJL MXNet JNA as the low level API
5. Use MXNet JavaCpp as the low level API
6. (Feel free to add more...)

Statistics

Scala package

Clojure package

@gigasquid @terrytangyuan @zachgk @frankfliu @aaronmarkham

[lanking520]

I would propose Option 3. and 4.

DJL is a new Java framework that build on top of any engines. It brings the ease for Java developers to have close to numpy experience in Java. It introduced an interface that defined Java to train and run inference on different ML/DL models.

In the engine layer, we implemented the MXNet specific engine that allows users to achieve most of the up-to-date functionalities:

MXNet specific

• deep-numpy ops (MXNet): numpy operators introduced in MXNet 1.6
• autograd (MXNet) : imperative training are supported by doing autogradient collections
• Block concept (MXNet Gluon): Java blocks for training and inference
• CachedOp (MXNet): new Symbolic Inference Engine

DJL

• Full training support: We support on Imperative/Symbolic training in Java.
• Memory collection: NDManager introduced to lively track and collect memory, ensured 100 hrs stable run in prod.
• MKLDNN: All ops and NDArray are built on top of MKLDNN acceleration
• ModelZoo: Model zoo is a central hub for all model storage, it contains model files, preprocessing and post processing logics.

Maintainance

• Better loading experience: All native MXNet files are uploaded to maven and users can easily switch different engine versions.
• JNA simplification: JNI layers are generated to save maintainance time and brought more logic into Java layer.

With the benefit listed above, I would recommend Option 3 to go for the DJL path since it already covered most up-to-date MXNet feature and supporting all different symbolic/imperative training/inference.

For Option 4: I am also thinking of bring our JNA layer back to MXNet so the community can build up their own Java/Scala frontend if they don't like DJL.

[terrytangyuan]

I propose option 1 and 2 since it took us a lot of efforts to bring MXNet to Scala originally and there are already adopters of Scala API in industries (some may not have been disclosed yet). But I am open to other options. Not familiar with DJL though but I assume @frankfliu and @lanking520 are the masters behind DJL.

[gigasquid]

@lanking520 thanks for the clarification above. A further question - How do you envision a current Scala MXNet user migrate their code? Is it going to be mostly reusable or is it going to be a complete rewrite for them?

[zachgk]

It is going to be closer to a complete rewrite. On the other hand, making a new Scala API would be imperative instead of symbolic and I think there are going to be a lot of operator changes to better match numpy in 2.0. I don't think the migration costs for a Scala 2.0 would be that much less anyway

For users who don't want a full rewrite, they can continue using an old release or whatever new releases we make on the v1.x branch.

[gigasquid]

For the Clojure package. It is a lot easier to interop with Java than with Scala - so if the the base is Java that everything is using - it will be better for Clojure.

[szha]

+1 for option 1 and 2. Also +1 for 4 as long as it doesn't add a dependency

My concerns on 3 and 4 are that DJL is a separate project which has its own release cycle. Having it to support MXNet's inference will cause delays as DJL upgrades to the latest version. This will also complicate the testing and validation.

Overall, I think a minimum set of API for at least inference is needed for MXNet JVM ecosystem users.

[leezu]

Another data point is that all of our Scala tests fail randomly with src/c_api/c_api_profile.cc:141: Check failed: !thread_profiling_data.calls_.empty():, so there seem to be some underlying issues.

#17067