Skip to content

Intel® Extension for OpenXLA* 0.4.0 Release

Latest
Latest
Compare
Choose a tag to compare
@Zantares Zantares released this 12 Aug 07:21
0.4.0
8c55fc3

Major Features

Intel® Extension for OpenXLA* is Intel optimized PyPI package to extend official OpenXLA on Intel GPU. It is based on PJRT plugin mechanism, which can seamlessly run JAX models on Intel® Data Center GPU Max Series and Intel® Data Center GPU Flex Series. This release contains following major features:

  • Jax Upgrade: Upgrade version to v0.4.26 and support the Compatibility of jax and jaxlib, it allows the Extension to support multiple different versions of jax. Please refer to <How are jax and jaxlib versioned?> for more version details between jax and jaxlib.
    intel-extension-for-openxla jaxlib jax
    0.4.0 0.4.26 >= 0.4.26, <= 0.4.27
  • Feature Support:
    • Support Float8 training and inference based on Keras 3.0. A new FP8 case is added to example.
    • Continue to improve jax native distributed scale-up collectives. A new distributed scale-up inference case Grok is added to example.
    • Support the FMHA backward fusion on Intel GPU.
  • Bug Fix:
    • Fix crash in jax native multi-process API.
    • Fix an accuracy error in dynamic slice fusion.
    • Fix known caveat crash related to Binary operations and SPMD multi-device parallelism API psum_scatter under same partial annotation.
    • Fix known caveat hang related to deadlock when working with Toolkit 2024.1.
    • Fix known caveat OOM related to deprecated API clear_backends.
  • Toolkit Support: Support Intel® oneAPI Base Toolkit 2024.2.
  • Driver Support: Support upgraded Driver LTS release 2350.63
  • OneDNN support: Support oneDNN v3.5.1.

Known Caveats

  • Some models show performance regression when working with Toolkit 2024.2. Recommend to use Toolkit 2024.1 if meet performance issues.
  • Multi-process API support is still experimental and may cause hang issues with collectives.

Breaking changes

  • Previous JAX v0.4.24 is no longer supported by this release. Please follow JAX change log to update the application if meets version errors. Please roll back the Extension version if want to use it with old JAX version.

Documents

Assets 2
Loading