kaushalya / jax_playground Goto Github PK

coremltools currently supports conversion from Tensorflow and PyTorch. As an alternative, we can convert JAX models to Tensorflow using jax2tf and then convert the generated Tensorflow model into coreml using coremltools.

Known issues:

• Only TF models a single concrete function can be converted to coreml.
• All TF/JAX operations are not supported by coremltools and Apple Silicon devices.

Device: M2 Max
Library versions
jax: 0.4.20
jaxlib: 0.4.20
jax-metal: 0.0.5

Decoder-only transformer with rotary positional embedding (RoPE) doesn't work on Apple Silicon device with jax-metal. The model generates gibberish text.
Looks like einsum and einops are not supported yet. It works fine when learned positional embeddings are used as it doesn't use einops.
On Apple M2 Max with jax-metal plugin:

Metal device set to: Apple M2 Max

systemMemory: 64.00 GB
maxCacheSize: 24.00 GB

Devices: [METAL(id=0)]
Using rope embeddings: True
Loading Tiny Shakespeare dataset...
Loaded Tiny Shakespeare dataset
Loading model parameters from:  /Users/kaushalya/code/jax_playground/artifacts/transformer_model:v3/transformer_epoch_100.pkl
Generating text...
fhzzzzzzzU?xfzzDzn.HLKxU-!zUzzBk?N'AEMPJzizhYk;VzD

The same model generates sensible output on CUDA devices.

Devices: [cuda(id=0), cuda(id=1)]
Using rope embeddings: True
Loading Tiny Shakespeare dataset...
Loaded Tiny Shakespeare dataset
Loading model parameters from:  models/transformer_rope_epoch_100.pkl
Generating text...
 you not to see your grace above and well,
Gaunt t

Paper: "RoFormer: Enhanced Transformer with Rotary Position Embedding" Arxiv, 2021